remove cpu prereqs (#3410)

This commit is contained in:
George Hotz
2024-02-15 13:45:06 +01:00
committed by GitHub
parent a40df14fef
commit 93eceef727
3 changed files with 4 additions and 4 deletions

View File

@@ -46,7 +46,7 @@ def _infer(model: EfficientNet, img, bs=1):
img = preprocess(img)
# run the net
if bs > 1: img = img.repeat(bs, axis=0)
out = model.forward(Tensor(img)).cpu()
out = model.forward(Tensor(img))
return _LABELS[np.argmax(out.numpy()[0])]
chicken_img = Image.open(pathlib.Path(__file__).parent / 'efficientnet/Chicken.jpg')

View File

@@ -464,7 +464,7 @@ class TestZeroShapeTensor(unittest.TestCase):
class TestTensorCreationDevice(unittest.TestCase):
# test auxiliary tensors are created on the same device
def test_one_hot(self):
y = Tensor([1, 2, 3]).to("CPU")
y = Tensor([1, 2, 3]).to("CLANG")
x = y.one_hot(10)
x.realize()

View File

@@ -168,7 +168,7 @@ class TestSafetensors(unittest.TestCase):
def helper_test_disk_tensor(fn, data, np_fxn, tinygrad_fxn=None):
if tinygrad_fxn is None: tinygrad_fxn = np_fxn
pathlib.Path(temp(fn)).unlink(missing_ok=True)
tinygrad_tensor = Tensor(data, device="CPU").to(f"disk:{temp(fn)}")
tinygrad_tensor = Tensor(data, device="CLANG").to(f"disk:{temp(fn)}")
numpy_arr = np.array(data)
tinygrad_fxn(tinygrad_tensor)
np_fxn(numpy_arr)
@@ -182,7 +182,7 @@ class TestDiskTensor(unittest.TestCase):
def test_write_ones(self):
pathlib.Path(temp("dt2")).unlink(missing_ok=True)
out = Tensor.ones(10, 10, device="CPU").contiguous()
out = Tensor.ones(10, 10, device="CLANG").contiguous()
outdisk = out.to(f"disk:{temp('dt2')}")
print(outdisk)
outdisk.realize()