Merge pull request #20 from Jeanselme/patch-2

Change fit return to match Sklearn format
autonlab · Dec 28, 2020 · e66f5db · e66f5db
2 parents 4093ba5 + 7b1b943
commit e66f5db
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/dsm/dsm_api.py b/dsm/dsm_api.py
@@ -116,6 +116,8 @@ def fit(self, x, t, e, vsize=0.15,
 
     self.torch_model = model.eval()
     self.fitted = True
+
+    return self
 
   def _prepocess_test_data(self, x):
     return torch.from_numpy(x)