fixed UT for pytorch 1.0

civodlu · Jan 29, 2020 · 6ee062d · 6ee062d
1 parent 7aecf07
commit 6ee062d
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/test_callback_learning_rate_finder.py b/tests/test_callback_learning_rate_finder.py
@@ -8,8 +8,8 @@
 
 
 def create_simple_regression(factor, nb_samples=100):
-    i = np.random.randn(nb_samples, 1)
-    o = i * factor
+    i = np.random.randn(nb_samples, 1).astype(np.float32)
+    o = i * np.float32(factor)
 
     datasets = collections.OrderedDict()
     sampler = trw.train.SamplerRandom(batch_size=256)