Use .to(device) instead of .cuda()

Author: tristandeleu

main.py

@@ -19,6 +19,7 @@
 LAMDA = 1
 LR = 2e-4
 
+DEVICE = torch.device('cuda') # torch.device('cpu')
 
 preproc_transform = transforms.Compose([
     transforms.ToTensor(),
@@ -39,15 +40,15 @@
     num_workers=NUM_WORKERS, pin_memory=True
 )
 
-model = AutoEncoder(INPUT_DIM, DIM, K).cuda()
+model = AutoEncoder(INPUT_DIM, DIM, K).to(DEVICE)
 opt = torch.optim.Adam(model.parameters(), lr=LR)
 
 
 def train():
     train_loss = []
     for batch_idx, (x, _) in enumerate(train_loader):
         start_time = time.time()
-        x = x.cuda()
+        x = x.to(DEVICE)
 
         opt.zero_grad()
 
@@ -85,7 +86,7 @@ def test():
     start_time = time.time()
     val_loss = []
     for batch_idx, (x, _) in enumerate(test_loader):
-        x = x.cuda()
+        x = x.to(DEVICE)
         x_tilde, z_e_x, z_q_x = model(x)
         loss_recons = F.mse_loss(x_tilde, x)
         loss_vq = F.mse_loss(z_q_x, z_e_x.detach())
@@ -100,7 +101,7 @@ def test():
 
 def generate_samples():
     x, _ = test_loader.__iter__().next()
-    x = x[:32].cuda()
+    x = x[:32].to(DEVICE)
     x_tilde, _, _ = model(x)
 
     x_cat = torch.cat([x, x_tilde], 0)

modules.py

@@ -190,7 +190,9 @@ def forward(self, x, label):
         return self.output_conv(x_h)
 
     def generate(self, label, shape=(8, 8), batch_size=64):
-        x = torch.zeros(batch_size, *shape).long().cuda()
+        param = next(self.parameters())
+        x = torch.zeros(batch_size, *shape).long()
+        x = x.to(param.device)
 
         for i in range(shape[0]):
             for j in range(shape[1]):

pixelcnn.py

@@ -22,6 +22,7 @@
 K = 512
 LR = 3e-4
 
+DEVICE = torch.device('cuda') # torch.device('cpu')
 
 preproc_transform = transforms.Compose([
     transforms.ToTensor(),
@@ -42,23 +43,23 @@
     num_workers=NUM_WORKERS, pin_memory=True
 )
 
-autoencoder = AutoEncoder(INPUT_DIM, VAE_DIM, K).cuda()
+autoencoder = AutoEncoder(INPUT_DIM, VAE_DIM, K).to(DEVICE)
 autoencoder.load_state_dict(
     torch.load('models/{}_autoencoder.pt'.format(DATASET))
 )
 autoencoder.eval()
 
-model = GatedPixelCNN(K, DIM, N_LAYERS).cuda()
-criterion = nn.CrossEntropyLoss().cuda()
+model = GatedPixelCNN(K, DIM, N_LAYERS).to(DEVICE)
+criterion = nn.CrossEntropyLoss().to(DEVICE)
 opt = torch.optim.Adam(model.parameters(), lr=LR)
 
 
 def train():
     train_loss = []
     for batch_idx, (x, label) in enumerate(train_loader):
         start_time = time.time()
-        x = x.cuda()
-        label = label.cuda()
+        x = x.to(DEVICE)
+        label = label.to(DEVICE)
 
         # Get the latent codes for image x
         latents, _ = autoencoder.encode(x)
@@ -93,8 +94,8 @@ def test():
     val_loss = []
     with torch.no_grad():
         for batch_idx, (x, label) in enumerate(test_loader):
-            x = x.cuda()
-            label = label.cuda()
+            x = x.to(DEVICE)
+            label = label.to(DEVICE)
 
             latents, _ = autoencoder.encode(x)
             logits = model(latents.detach(), label)
@@ -114,7 +115,7 @@ def test():
 
 def generate_samples():
     label = torch.arange(10).expand(10, 10).contiguous().view(-1)
-    label = label.long().cuda()
+    label = label.long().to(DEVICE)
 
     latents = model.generate(label, shape=LATENT_SHAPE, batch_size=100)
     x_tilde, _ = autoencoder.decode(latents)
@@ -129,7 +130,7 @@ def generate_samples():
 
 def generate_reconstructions():
     x, _ = test_loader.__iter__().next()
-    x = x[:32].cuda()
+    x = x[:32].to(DEVICE)
 
     latents, _ = autoencoder.encode(x)
     x_tilde, _ = autoencoder.decode(latents)