pixart special

main.py

@@ -202,8 +202,12 @@ def execute_task(args, trainer, device, dtype, shape, enable_grad, settings, pip
     if args.task == "single":
         # Attempt to move the model to GPU if model is not Flux
         if args.model != "flux":
-            if pipe.device != torch.device('cuda'):
-                pipe.to(device, dtype)
+            if args.model != "pixart":
+                if pipe.device != torch.device('cuda'):
+                    pipe.to(device, dtype)
+            else: 
+                if pipe.device != torch.device('cuda'):
+                    pipe.to(device)
         else:
             print(f"PIPE:{pipe}")