def test_default_predict_fn_gpu_gpu(tensor): tensor = tensor.cuda() model = DummyModel().cuda() prediction = default_predict_fn(tensor, model) assert torch.equal(model(tensor), prediction) assert prediction.is_cuda is True
def test_default_predict_fn(tensor): model = DummyModel() prediction = default_predict_fn(tensor, model) assert torch.equal(model(Variable(tensor)), prediction) assert prediction.is_cuda == torch.cuda.is_available()
def test_default_predict_fn_cpu_cpu(tensor): prediction = default_predict_fn(tensor.cpu(), DummyModel().cpu()) model = DummyModel().to(device) assert torch.equal(model(Variable(tensor)), prediction) assert prediction.is_cuda == torch.cuda.is_available()