Explicitly disable nativize tensor

thatguy11325 · Jun 29, 2024 · 9435fba · 9435fba
1 parent 06fdffa
commit 9435fba
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/pokemonred_puffer/policies/multi_convolutional.py b/pokemonred_puffer/policies/multi_convolutional.py
@@ -8,6 +8,8 @@
 from pokemonred_puffer.data.items import Items
 from pokemonred_puffer.environment import PIXEL_VALUES
 
+pufferlib.pytorch.nativize_tensor = torch.compiler.disable(pufferlib.pytorch.nativize_tensor)
+
 
 # Because torch.nn.functional.one_hot cannot be traced by torch as of 2.2.0
 def one_hot(tensor, num_classes):