Update modularStarEncoder.py

modularStarEncoder.py

@@ -205,15 +205,15 @@ def get_pooling_mask(
     repeated_idx = idx.unsqueeze(1).repeat(1, input_ids.size(1))
 
     DEVICE = input_ids.get_device()
-    print(DEVICE)
 
     if DEVICE<0:
         DEVICE = "cpu"
     ranges = torch.arange(input_ids.size(1)).repeat(input_ids.size(0), 1)
+    ranges.to(DEVICE)
+    repeated_idx.to(DEVICE)
 
-    print(repeated_idx.get_device(),ranges.get_device())
     pooling_mask = (repeated_idx <= ranges).long()
-    pooling_mask.to(DEVICE)
+
 
     return pooling_mask