def test_prediction_output(self):
model = SimpleModel()
dp = DataProcessor(model=model)
self.assertFalse(model.fc.weight.is_cuda)
res = dp.predict({'data': torch.rand(1, 3)})
self.assertIs(type(res), torch.Tensor)
model = NonStandardIOModel()
dp = DataProcessor(model=model)
self.assertFalse(model.fc.weight.is_cuda)
res = dp.predict({'data': {'data1': torch.rand(1, 3), 'data2': torch.rand(1, 3)}})
self.assertIs(type(res), dict)
self.assertIn('res1', res)
self.assertIs(type(res['res1']), torch.Tensor)
self.assertIn('res2', res)
self.assertIs(type(res['res2']), torch.Tensor)
def test_predict(self):
model = SimpleModel().train()
dp = DataProcessor(model=model)
self.assertFalse(model.fc.weight.is_cuda)
self.assertTrue(model.training)
res = dp.predict({'data': torch.rand(1, 3)})
self.assertFalse(model.training)
self.assertFalse(res.requires_grad)
self.assertIsNone(res.grad)