modified: tests/test_dsm.py

autonlab · Oct 28, 2020 · d1c2197 · d1c2197
1 parent a3f0fc8
commit d1c2197
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/tests/test_dsm.py b/tests/test_dsm.py
@@ -1,6 +1,7 @@
 import unittest
 
 from dsm import DeepSurvivalMachines
+from dsm.dsm_torch import DeepSurvivalMachinesTorch
 from dsm import datasets
 
 import numpy as np
@@ -20,7 +21,7 @@ def test_dsm(self):
     self.assertEqual(e.shape, (9105,))
 
     model = DeepSurvivalMachines()
-    self.assertIsInstance(model, dsm.dsm_api.DeepSurvivalMachines)
+    self.assertIsInstance(model, DeepSurvivalMachines)
     model.fit(x, t, e, iters=10)
     self.assertIsInstance(model.torch_model,
-                          dsm.dsm_torch.DeepSurvivalMachinesTorch)
+                          DeepSurvivalMachinesTorch)