mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-09-03 16:45:51 +00:00
update embedding test to not use dtypes.long [pr] (#9556)
This commit is contained in:
@@ -307,12 +307,11 @@ class TestMultiTensor(unittest.TestCase):
|
||||
lr_sched = OneCycleLR(optim, max_lr=0.1, pct_start=0.1, div_factor=100, final_div_factor=0.1, total_steps=10)
|
||||
lr_sched.step()
|
||||
|
||||
@unittest.skipUnless(is_dtype_supported(dtypes.long), f"long dtype not supported on {Device.DEFAULT}")
|
||||
def test_embedding(self):
|
||||
B, T, embed_size, vocab_size = 4, 10, 20, 28
|
||||
|
||||
layer = nn.Embedding(vocab_size, embed_size)
|
||||
x = Tensor(np.random.randint(0, vocab_size, (B, T)))
|
||||
x = Tensor(np.random.randint(0, vocab_size, (B, T), dtype=np.int32))
|
||||
z = layer(x)
|
||||
|
||||
layer_sharded = nn.Embedding(vocab_size, embed_size)
|
||||
|
||||
Reference in New Issue
Block a user