Update modularStarEncoder.py
Browse files- modularStarEncoder.py +4 -5
modularStarEncoder.py
CHANGED
@@ -209,13 +209,12 @@ def get_pooling_mask(
|
|
209 |
if DEVICE<0:
|
210 |
DEVICE = "cpu"
|
211 |
ranges = torch.arange(input_ids.size(1)).repeat(input_ids.size(0), 1)
|
212 |
-
ranges.to(DEVICE)
|
213 |
-
|
214 |
-
|
215 |
-
print(ranges.get_device(),repeated_idx.get_device())
|
216 |
|
217 |
pooling_mask = (repeated_idx <= ranges).long()
|
218 |
-
|
219 |
|
220 |
return pooling_mask
|
221 |
|
|
|
209 |
if DEVICE<0:
|
210 |
DEVICE = "cpu"
|
211 |
ranges = torch.arange(input_ids.size(1)).repeat(input_ids.size(0), 1)
|
212 |
+
ranges = ranges.to(DEVICE)
|
213 |
+
|
214 |
+
|
|
|
215 |
|
216 |
pooling_mask = (repeated_idx <= ranges).long()
|
217 |
+
|
218 |
|
219 |
return pooling_mask
|
220 |
|