sandy898 commited on
Commit
f33ea10
·
1 Parent(s): d597de7

Fix tensor dtype mismatch in ranker

Browse files
Files changed (1) hide show
  1. app/services/recommender.py +23 -4
app/services/recommender.py CHANGED
@@ -140,10 +140,29 @@ class RecommenderService:
140
  autocast_ctx = torch.autocast(device_type=mgr.device.type) if mgr.device.type == 'cuda' else contextlib.nullcontext()
141
 
142
  with torch.no_grad(), autocast_ctx:
143
- t_tab = torch.tensor(tab_feats, device=mgr.device).unsqueeze(0)
144
- t_f = torch.tensor(franchise_ids, device=mgr.device).unsqueeze(0)
145
- t_cands = torch.tensor(padded_cands, device=mgr.device).unsqueeze(0)
146
- t_mask = torch.tensor(valid_mask, device=mgr.device).unsqueeze(0)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
147
 
148
  logits = mgr.ranker(t_tab, t_f, t_cands, t_mask).squeeze(0)
149
  logits = logits.masked_fill(~torch.tensor(valid_mask, device=mgr.device), -float('inf'))
 
140
  autocast_ctx = torch.autocast(device_type=mgr.device.type) if mgr.device.type == 'cuda' else contextlib.nullcontext()
141
 
142
  with torch.no_grad(), autocast_ctx:
143
+ t_tab = torch.tensor(
144
+ tab_feats,
145
+ dtype=torch.float32,
146
+ device=mgr.device
147
+ ).unsqueeze(0)
148
+
149
+ t_f = torch.tensor(
150
+ franchise_ids,
151
+ dtype=torch.long,
152
+ device=mgr.device
153
+ ).unsqueeze(0)
154
+
155
+ t_cands = torch.tensor(
156
+ padded_cands,
157
+ dtype=torch.long,
158
+ device=mgr.device
159
+ ).unsqueeze(0)
160
+
161
+ t_mask = torch.tensor(
162
+ valid_mask,
163
+ dtype=torch.bool,
164
+ device=mgr.device
165
+ ).unsqueeze(0)
166
 
167
  logits = mgr.ranker(t_tab, t_f, t_cands, t_mask).squeeze(0)
168
  logits = logits.masked_fill(~torch.tensor(valid_mask, device=mgr.device), -float('inf'))