Fix pytorch storage warnings, fixes #80 (#88)

we seriously don't care what type of storage we get, pytorch sucks
RWKV · Jun 3, 2023 · fb6708b · fb6708b
1 parent 3f8bb2c
commit fb6708b
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/rwkv/rwkv_cpp_model.py b/rwkv/rwkv_cpp_model.py
@@ -84,7 +84,7 @@ def validate_buffer(buf: torch.Tensor, name: str, size: int) -> None:
         if state_in is not None:
             validate_buffer(state_in, 'state_in', self._state_buffer_element_count)
 
-            state_in_ptr = state_in.untyped_storage().data_ptr()
+            state_in_ptr = state_in.data_ptr()
         else:
             state_in_ptr = 0
 
@@ -102,8 +102,8 @@ def validate_buffer(buf: torch.Tensor, name: str, size: int) -> None:
             self._ctx,
             token,
             state_in_ptr,
-            state_out.untyped_storage().data_ptr(),
-            logits_out.untyped_storage().data_ptr()
+            state_out.data_ptr(),
+            logits_out.data_ptr()
         )
 
         return logits_out, state_out