From d1c21971e4b143cd5768d2dd9c2dc0d0686bd327 Mon Sep 17 00:00:00 2001 From: Chirag Nagpal Date: Wed, 28 Oct 2020 19:19:42 -0400 Subject: [PATCH] modified: tests/test_dsm.py --- tests/test_dsm.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/test_dsm.py b/tests/test_dsm.py index b684946..fe3fe43 100644 --- 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)