update

app.py

@@ -175,8 +175,10 @@ def main(args):
     # denoise_model.load_state_dict(
     #     dist_util.load_state_dict(args.ddpm_model_path, map_location="cpu"))
     denoise_model.to(dist_util.dev())
-    if args.use_fp16:
-        denoise_model.convert_to_fp16()
+    denoise_model = denoise_model.to(th.bfloat16)
+    auto_encoder = auto_encoder.to(th.bfloat16)
+    # if args.use_fp16:
+    #     denoise_model.convert_to_fp16()
     denoise_model.eval()
 
     # * auto-encoder reconstruction model