compile disable nativize tensor

thatguy11325 · Jul 11, 2024 · 4046ea9 · 4046ea9
1 parent 85fbb1c
commit 4046ea9
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):