hotfix for pytorch 0.4.1

masenov · Nov 15, 2018 · e467dbb · e467dbb
1 parent 220c634
commit e467dbb
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/examples/odenet_mnist.py b/examples/odenet_mnist.py
@@ -105,6 +105,7 @@ def __init__(self, odefunc):
         self.integration_time = torch.tensor([0, 1]).float()
 
     def forward(self, x):
+        self.integration_time = self.integration_time.type_as(x)
         out = odeint(self.odefunc, x, self.integration_time, rtol=args.tol, atol=args.tol)
         return out[1]