update embedding test to not use dtypes.long [pr] (#9556)

This commit is contained in:
chenyu
2025-03-23 21:33:38 -04:00
committed by GitHub
parent c965f4c20b
commit ba41076e94
2 changed files with 5 additions and 8 deletions
+1 -2
View File
@@ -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)