fix float64 warning on training
MODIFY
test/test_mnist.py
test/test_mnist.py
@@ -5,7 +5,7 @@ import numpy as np
from tinygrad.tensor import Tensor
import tinygrad.optim as optim
from extra.training import train, evaluate
-from extra.utils import fetch, get_parameters
+from extra.utils import get_parameters
from datasets import fetch_mnist
# load the mnist dataset

MODIFY
tinygrad/ops_cpu.py
tinygrad/ops_cpu.py
@@ -58,7 +58,7 @@ class Max(Function):
shape = [1 if axis is None or i in axis else input.shape[i] for i in range(len(input.shape))]
ret2 = (input==ret.reshape(shape))
div = ret2.sum(axis=tuple(axis), keepdims=True) if axis is not None else ret2.sum()
- return ret2*grad_output.reshape(shape)/div
+ return ret2*grad_output.reshape(shape)/div.type(input.dtype)
# ************* binary ops *************

MODIFY
tinygrad/tensor.py
tinygrad/tensor.py
@@ -185,6 +185,8 @@ class Tensor:
return np.amax(x, *args, **kwargs)
def permute(x, order):
return x.transpose(order)
+ def type(x, tt):
+ return x.astype(tt)
data = data.view(CPUBuffer)
if device == Device.GPU:

fix float64 warning on training