amend

pytorch · Jun 24, 2024 · 91e62d1 · 91e62d1
1 parent 554fb15
commit 91e62d1
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tensordict/nn/storage.py b/tensordict/nn/storage.py
@@ -143,7 +143,7 @@ def __init__(
     ):
         super().__init__()
         self.convert_to_binary = convert_to_binary
-        self.bases = 2 ** torch.arange(num_bits - 1, -1, -1).to(device, dtype)
+        self.bases = 2 ** torch.arange(num_bits - 1, -1, -1, device=device, dtype=dtype)
         self.num_bits = num_bits
         self.zero_tensor = torch.zeros((1,))