Bug: torch.jit.load cannot load from IOBuffer

jit
1

Taking code straight from the documentation on this page:

import torch
import io

class MyModule(torch.nn.Module):
    def forward(self, x):
        return x + 10

m = torch.jit.script(MyModule())

# Save to io.BytesIO buffer
buffer = io.BytesIO()
torch.jit.save(m, buffer)

If we then run:

torch.jit.load(buffer)

We see:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Library/Frameworks/Python.framework/Versions/3.9/lib/python3.9/site-packages/torch/jit/_serialization.py", line 163, in load
    cpp_module = torch._C.import_ir_module_from_buffer(
RuntimeError: PytorchStreamReader failed reading zip archive: not a ZIP archive

This is on Mac and Linux, Python3.9, torch 1.9.0.

2

you probably wanto do a buffer.seek(0) to reset the cursor

1 Like