From 9b27ba650b5d7d8742e14f9aef8a4e3765cd7655 Mon Sep 17 00:00:00 2001
From: George Hotz <geohot@gmail.com>
Date: Tue, 7 Jun 2022 10:06:48 -0700
Subject: [PATCH] load new torch files

---
 extra/utils.py | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/extra/utils.py b/extra/utils.py
index 7ce050531d..18eb2266ed 100644
--- a/extra/utils.py
+++ b/extra/utils.py
@@ -76,7 +76,21 @@ def my_unpickle(fb0):
 
   return MyPickle(fb0).load(), key_prelookup
 
+def fake_torch_load_zipped(b0):
+  import io, zipfile
+  with zipfile.ZipFile(io.BytesIO(b0), 'r') as myzip:
+    with myzip.open('archive/data.pkl') as myfile:
+      ret = my_unpickle(myfile)
+    for k,v in ret[1].items():
+      with myzip.open(f'archive/data/{k}') as myfile:
+        v[2][:] = np.frombuffer(myfile.read(), v[2].dtype)
+        v[2].shape = v[3]
+  return ret[0]
+
 def fake_torch_load(b0):
+  if b0[0:2] == b"\x50\x4b":
+    return fake_torch_load_zipped(b0)
+
   import io
   import struct