Update train.py

train.py

@@ -7,13 +7,13 @@ def train(prompt, modules, freeze_modules, iterations, negative_guidance, lr, sa
   
     nsteps = 50
 
-    diffuser = StableDiffuser(scheduler='DDIM').to('cuda:1')
+    diffuser = StableDiffuser(scheduler='DDIM').to('cuda')
     diffuser.train()
 
 
 
 
-    finetuner = FineTunedModel(diffuser, modules, frozen_modules=freeze_modules).to('cuda:0')
+    finetuner = FineTunedModel(diffuser, modules, frozen_modules=freeze_modules)
 
     optimizer = torch.optim.Adam(finetuner.parameters(), lr=lr)
     criteria = torch.nn.MSELoss()
@@ -58,8 +58,8 @@ def train(prompt, modules, freeze_modules, iterations, negative_guidance, lr, sa
 
             iteration = int(iteration / nsteps * 1000)
             
-            positive_latents = diffuser.predict_noise(iteration, latents_steps[0], positive_text_embeddings, guidance_scale=1).to('cuda:0')
-            neutral_latents = diffuser.predict_noise(iteration, latents_steps[0], neutral_text_embeddings, guidance_scale=1).to('cuda:0')
+            positive_latents = diffuser.predict_noise(iteration, latents_steps[0], positive_text_embeddings, guidance_scale=1)
+            neutral_latents = diffuser.predict_noise(iteration, latents_steps[0], neutral_text_embeddings, guidance_scale=1)
 
         with finetuner:
             negative_latents = diffuser.predict_noise(iteration, latents_steps[0], positive_text_embeddings, guidance_scale=1)