Update osediff_sd3.py

osediff_sd3.py

@@ -392,7 +392,7 @@ class OSEDiff_SD3_GEN(torch.nn.Module):
         # Add lora to transformer
         print('Adding Lora to OSEDiff_SD3_GEN')
         self.transformer_gen = copy.deepcopy(self.model.transformer)
-        self.transformer_gen.to('cuda:1')
+        self.transformer_gen.to('cuda')
         # self.transformer_gen = self.transformer_gen.float()
         
         self.transformer_gen.requires_grad_(False)
@@ -468,7 +468,7 @@ class OSEDiff_SD3_REG(torch.nn.Module):
         # Add lora to transformer
         print('Adding Lora to OSEDiff_SD3_REG')
         self.transformer_reg = copy.deepcopy(self.transformer_org)
-        self.transformer_reg.to('cuda:1')
+        self.transformer_reg.to('cuda')
         
         self.transformer_reg.requires_grad_(False)
         self.transformer_reg.train()