modularStarEncoder
/

ModularStarEncoder-finetuned-4

Feature Extraction

ModularStarEncoder

Model card Files Files and versions Community

andreagurioli1995 commited on Feb 20

Commit

739f9ab

·

verified ·

1 Parent(s): 8512abe

Update modularStarEncoder.py

Files changed (1) hide show

modularStarEncoder.py +2 -1

modularStarEncoder.py CHANGED Viewed

@@ -208,9 +208,10 @@ def get_pooling_mask(
     if DEVICE<0:
         DEVICE = "cpu"
     ranges = torch.arange(input_ids.size(1)).repeat(input_ids.size(0), 1)
-    ranges.to(DEVICE)
     pooling_mask = (repeated_idx <= ranges).long()
     return pooling_mask

     if DEVICE<0:
         DEVICE = "cpu"
     ranges = torch.arange(input_ids.size(1)).repeat(input_ids.size(0), 1)
     pooling_mask = (repeated_idx <= ranges).long()
+    pooling_mask.to(DEVICE)
     return pooling_mask