Method create_model

ML/tests/test_model.py:55–57

(self, num_classes)

        print("=" * 60)
    
    def create_model(self, num_classes):
        model = ResNet18(num_classes=num_classes)
        return model.to(self.device)
    
    def predict_image(self, image_tensor):
        with torch.no_grad():