Update modularStarEncoder.py

modularStarEncoder.py

@@ -211,9 +211,11 @@ def get_pooling_mask(
     ranges = torch.arange(input_ids.size(1)).repeat(input_ids.size(0), 1)
     ranges.to(DEVICE)
     repeated_idx.to(DEVICE)
+    print(ranges,repeated_idx)
+    print(ranges.get_device(),repeated_idx.get_device())
 
     pooling_mask = (repeated_idx <= ranges).long()
-
+    print(pooling_mask.get_device())
 
     return pooling_mask