fix: use untyped storage to stop warning

aphp · Nov 21, 2024 · d4cde5d · d4cde5d
1 parent 3f02fb1
commit d4cde5d
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/changelog.md b/changelog.md
@@ -1,5 +1,9 @@
 # Changelog
 
+## Unreleased
+
+- Fix `storage` torch warning
+
 ## v0.3.5
 
 - Support hashing the `folded_tensor.length` field (via a UserList), which is convenient for caching

diff --git a/foldedtensor/__init__.py b/foldedtensor/__init__.py
@@ -428,12 +428,12 @@ def reduce_foldedtensor(self: FoldedTensor):
             self.data_dims,
             self.full_names,
             self.indexer.clone()
-            if self.indexer.is_shared() and self.indexer.storage().is_cuda
+            if self.indexer.is_shared() and self.indexer.untyped_storage().is_cuda
             else self.indexer,
             None
             if self._mask is not None
             and self._mask.is_shared()
-            and self._mask.storage().is_cuda
+            and self._mask.untyped_storage().is_cuda
             else self._mask,
         ),
     )