Feature Extraction
Transformers
Safetensors
ModularStarEncoder
custom_code
andreagurioli1995 commited on
Commit
009c4e5
·
verified ·
1 Parent(s): 7dc1938

Update modularStarEncoder.py

Browse files
Files changed (1) hide show
  1. 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
- repeated_idx.to(DEVICE)
214
- print(ranges,repeated_idx)
215
- print(ranges.get_device(),repeated_idx.get_device())
216
 
217
  pooling_mask = (repeated_idx <= ranges).long()
218
- print(pooling_mask.get_device())
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