fix: add ModuleList init and sync model device with input data #21

models/s3rec.py

@@ -26,6 +26,12 @@ def _init_weights(self):
         for child in self.children():
             if isinstance(child, nn.Embedding):
                 nn.init.xavier_uniform_(child.weight)
+            elif isinstance(child, nn.ModuleList): # nn.Linear):
+                for sub_child in child.children():
+                    if not isinstance(sub_child, nn.MultiheadAttention):
+                        nn.init.xavier_uniform_(sub_child.weight)
+            else:
+                logger.info(f"other type: {child} / {type(child)}")
 
     def _embedding_layer(self, X):
         return self.item_embedding(X) + self.positional_encoding

trainers/s3rec_trainer.py

@@ -21,7 +21,7 @@
 class S3RecTrainer(BaseTrainer):
     def __init__(self, cfg: DictConfig, num_items: int, num_users: int, item2attributes, attributes_count: int) -> None:
         super().__init__(cfg)
-        self.model = S3Rec(self.cfg, num_items, num_users, attributes_count)
+        self.model = S3Rec(self.cfg, num_items, num_users, attributes_count).to(self.device)
         self.optimizer: Optimizer = self._optimizer(self.cfg.optimizer, self.model, self.cfg.lr)
         self.loss = self._loss()
 
@@ -145,7 +145,7 @@ def _generate_target_and_top_k_recommendation(self, scores: Tensor, pos_item: Te
 
         # create item index information
         scores_idx = np.zeros_like(scores.cpu().detach().numpy())
-        scores_idx[:, 0] = pos_item
+        scores_idx[:, 0] = pos_item.cpu().detach()
 
         # sort topK probs and find their indexes
         sorted_indices = np.argsort(-scores.cpu().detach().numpy(), axis=1)[:, :self.cfg.top_n]