def test_gpu(self, coef, x, expected): device = "cuda" if torch.cuda.is_available() else "cpu" x = torch.as_tensor(x, dtype=torch.float, device=device) x.requires_grad = True coef = torch.as_tensor(coef, dtype=torch.float, device=device) coef.requires_grad = True result = polyval(coef, x) if coef.shape[0] > 0: # empty coef doesn't have grad result.mean().backward() np.testing.assert_allclose(coef.grad.shape, coef.shape) np.testing.assert_allclose(result.cpu().detach().numpy(), expected)
def test_floats(self, coef, x, expected): result = polyval(coef, x) np.testing.assert_allclose(result.cpu().numpy(), expected)