model touchups

This commit is contained in:
George Hotz
2021-11-30 11:13:34 -05:00
parent 7d7e2b690d
commit 46bbbcf7f0
5 changed files with 14 additions and 15 deletions

View File

@@ -42,7 +42,7 @@ class TestEfficientNet(unittest.TestCase):
def test_chicken(self):
chicken_img = Image.open(pathlib.Path(__file__).parent / 'efficientnet/Chicken.jpg')
model = EfficientNet(number=0)
model.load_weights_from_torch()
model.load_from_pretrained()
label = _infer(model, chicken_img)
self.assertEqual(label, "hen", f"Expected hen but got {label} for number=0")