Update modularStarEncoder.py

modularStarEncoder.py

@@ -205,11 +205,13 @@ 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)
 
-
+    print(repeated_idx.get_device(),ranges.get_device())
     pooling_mask = (repeated_idx <= ranges).long()
     pooling_mask.to(DEVICE)