Update models/unet.py

models/unet.py

@@ -759,24 +759,32 @@ class CondUnet1D(nn.Module):
         cond,
         cond_indices,
     ):
+        self.time_mlp = self.time_mlp.cuda()
         temb = self.time_mlp(t)
 
         h = []
         for block1, block2, downsample in self.downs:
+            block1 = block1.cuda()
+            block2 = block2.cuda()
             x = block1(x, temb, cond, cond_indices)
             x = block2(x, temb, cond, cond_indices)
             h.append(x)
             x = downsample(x)
 
+        self.mid_block1 = self.mid_block1.cuda()
+        self.mid_block2 = self.mid_block2.cuda()
         x = self.mid_block1(x, temb, cond, cond_indices)
         x = self.mid_block2(x, temb, cond, cond_indices)
 
         for upsample, block1, block2 in self.ups:
             x = upsample(x)
             x = torch.cat((x, h.pop()), dim=1)
+            block1 = block1.cuda()
+            block2 = block2.cuda()
             x = block1(x, temb, cond, cond_indices)
             x = block2(x, temb, cond, cond_indices)
 
+        self.final_conv = self.final_conv.cuda()
         x = self.final_conv(x)
         return x