fix bugs in parameter initialization in RMTPP

ant-research · May 8, 2024 · 5ab9bcf · 5ab9bcf
1 parent 7b7659a
commit 5ab9bcf
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/easy_tpp/model/torch_model/torch_rmtpp.py b/easy_tpp/model/torch_model/torch_rmtpp.py
@@ -24,8 +24,8 @@ def __init__(self, model_config):
 
         self.layer_hidden = nn.Linear(self.hidden_size, self.num_event_types)
 
-        self.factor_intensity_base = torch.empty([1, 1, self.num_event_types], device=self.device)
-        self.factor_intensity_current_influence = torch.empty([1, 1, self.num_event_types], device=self.device)
+        self.factor_intensity_base = torch.nn.Parameter(torch.empty([1, 1, self.num_event_types], device=self.device))
+        self.factor_intensity_current_influence = torch.nn.Parameter(torch.empty([1, 1, self.num_event_types], device=self.device))
 
         nn.init.xavier_normal_(self.factor_intensity_base)
         nn.init.xavier_normal_(self.factor_intensity_current_influence)