Update mar.py

mar.py

@@ -14,8 +14,8 @@ from diffloss import DiffLoss
 
 
 def mask_by_order(mask_len, order, bsz, seq_len):
-    masking = torch.zeros(bsz, seq_len)
-    masking = torch.scatter(masking, dim=-1, index=order[:, :mask_len.long()], src=torch.ones(bsz, seq_len)).bool()
+    masking = torch.zeros(bsz, seq_len).to(device)
+    masking = torch.scatter(masking, dim=-1, index=order[:, :mask_len.long()], src=torch.ones(bsz, seq_len).to(device)).bool()
     return masking
 
 
@@ -156,7 +156,7 @@ class MAR(nn.Module):
             order = np.array(list(range(self.seq_len)))
             np.random.shuffle(order)
             orders.append(order)
-        orders = torch.Tensor(np.array(orders)).long()
+        orders = torch.Tensor(np.array(orders)).to(device).long()
         return orders
 
     def random_masking(self, x, orders):
@@ -180,7 +180,7 @@ class MAR(nn.Module):
         # random drop class embedding during training
         if self.training:
             drop_latent_mask = torch.rand(bsz) < self.label_drop_prob
-            drop_latent_mask = drop_latent_mask.unsqueeze(-1).to(x.dtype)
+            drop_latent_mask = drop_latent_mask.unsqueeze(-1).to(device).to(x.dtype)
             class_embedding = drop_latent_mask * self.fake_latent + (1 - drop_latent_mask) * class_embedding
 
         x[:, :self.buffer_size] = class_embedding.unsqueeze(1)
@@ -262,8 +262,8 @@ class MAR(nn.Module):
     def sample_tokens(self, bsz, num_iter=64, cfg=1.0, cfg_schedule="linear", labels=None, temperature=1.0, progress=False):
 
         # init and sample generation orders
-        mask = torch.ones(bsz, self.seq_len)
-        tokens = torch.zeros(bsz, self.seq_len, self.token_embed_dim)
+        mask = torch.ones(bsz, self.seq_len).to(device)
+        tokens = torch.zeros(bsz, self.seq_len, self.token_embed_dim).to(device)
         orders = self.sample_orders(bsz)
 
         indices = list(range(num_iter))
@@ -291,10 +291,10 @@ class MAR(nn.Module):
 
             # mask ratio for the next round, following MaskGIT and MAGE.
             mask_ratio = np.cos(math.pi / 2. * (step + 1) / num_iter)
-            mask_len = torch.Tensor([np.floor(self.seq_len * mask_ratio)])
+            mask_len = torch.Tensor([np.floor(self.seq_len * mask_ratio)]).to(device)
 
             # masks out at least one for the next iteration
-            mask_len = torch.maximum(torch.Tensor([1]),
+            mask_len = torch.maximum(torch.Tensor([1]).to(device),
                                      torch.minimum(torch.sum(mask, dim=-1, keepdims=True) - 1, mask_len))
 
             # get masking for next iteration and locations to be predicted in this iteration