Co-authored-by: Brayden Zhong <brayden@radixark.ai>
This commit is contained in:
Brayden Zhong
2026-06-28 23:46:18 -07:00
committed by GitHub
co-authored by Brayden Zhong
parent c0d45cda29
commit c7b9b92d9a
+2 -2
View File
@@ -321,8 +321,8 @@ class TestDSAIndexer(CustomTestCase):
params.update(kwargs)
torch.set_default_dtype(self.dtype)
indexer = Indexer(**params)
# Move indexer to CUDA device
with torch.device(self.device):
indexer = Indexer(**params)
indexer = indexer.to(device=self.device)
# Convert linear layer weights to bfloat16 (but preserve LayerNorm's float32