Update models/unet.py

models/unet.py

@@ -273,6 +273,7 @@ class Downsample1d(nn.Module):
         self.conv = nn.Conv1d(dim, dim, 3, 2, 1)
 
     def forward(self, x):
+        self.conv = self.conv.cuda()
         return self.conv(x)
 
 
@@ -283,6 +284,7 @@ class Upsample1d(nn.Module):
         self.conv = nn.ConvTranspose1d(dim_in, dim_out, 4, 2, 1)
 
     def forward(self, x):
+        self.conv = self.conv.cuda()
         return self.conv(x)