From f95658bc3e1eb30966b308fae9dd58edbc5b9d74 Mon Sep 17 00:00:00 2001 From: George Hotz Date: Sun, 5 May 2024 10:14:03 -0700 Subject: [PATCH] hotfix: pickle jit works if you delete the function --- test/test_pickle.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/test/test_pickle.py b/test/test_pickle.py index 7833bfe6af..7955b17aa5 100644 --- a/test/test_pickle.py +++ b/test/test_pickle.py @@ -33,16 +33,15 @@ class TestPickle(unittest.TestCase): t2:Tensor = pickle.loads(st) np.testing.assert_equal(t.numpy(), t2.numpy()) - @unittest.expectedFailure def test_pickle_jit(self): @TinyJit def add(a, b): return a+b+1 for _ in range(3): add(Tensor.rand(10, 10), Tensor.rand(10, 10)) - #import dill - #with dill.detect.trace(): dill.dumps(add) + del add.fxn # pickling the JIT requires the function to be deleted st = pickle.dumps(add) - add_fxn = pickle.loads(st) + del add + add_fxn = pickle.loads(st) x = Tensor.ones(10, 10).contiguous().realize() y = Tensor.ones(10, 10).contiguous().realize() print("post jit")