Update mar.py

mar.py

@@ -12,7 +12,7 @@ from timm.models.vision_transformer import Block
 
 from .diffloss import DiffLoss
 
-
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 def mask_by_order(mask_len, order, bsz, seq_len):
     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()