this fix the gc
@@ -46,6 +46,10 @@ class TestEfficientNet(unittest.TestCase):
cls.model = EfficientNet(number=0)
cls.model.load_from_pretrained()
+ @classmethod
+ def tearDownClass(cls):
+ del cls.model
+
def test_chicken(self):
label = _infer(self.model, chicken_img)
self.assertEqual(label, "hen")
@@ -60,6 +64,10 @@ class TestViT(unittest.TestCase):
cls.model = ViT()
cls.model.load_from_pretrained()
+ @classmethod
+ def tearDownClass(cls):
+ del cls.model
+
def test_chicken(self):
label = _infer(self.model, chicken_img)
self.assertEqual(label, "cock")
