fix bug #261

shenweichen · Feb 11, 2023 · 5d03ead · 5d03ead
1 parent f685425
commit 5d03ead
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/deepctr_torch/models/basemodel.py b/deepctr_torch/models/basemodel.py
@@ -76,7 +76,7 @@ def forward(self, X, sparse_feat_refine_weight=None):
 
         sparse_embedding_list += varlen_embedding_list
 
-        linear_logit = torch.zeros([X.shape[0], 1]).to(self.device)
+        linear_logit = torch.zeros([X.shape[0], 1]).to(self.weight.device)
         if len(sparse_embedding_list) > 0:
             sparse_embedding_cat = torch.cat(sparse_embedding_list, dim=-1)
             if sparse_feat_refine_weight is not None: