`torch.save` and `torch.load` fail to roundtrip with `io.BytesIO`

[jatentaki]

🐛 Bug

Docs for torch.save suggest that saving to io.BytesIO should work, but trying to torch.load such a buffer causes EOFError: Ran out of input.

To Reproduce

Check the colab: https://colab.research.google.com/drive/1hY-8r5CWJ-tpq8wC-2MsR2vTAkKPTbov?usp=sharing

Copied below for clarity:

import io
import torch
import torch.nn as nn

example_model = nn.Sequential(
    nn.Linear(10, 10),
    nn.ReLU(),
    nn.Linear(10, 10),
)

buffer = io.BytesIO()
torch.save(example_model.state_dict(), buffer)
loaded = torch.load(buffer)

EOFError                                  Traceback (most recent call last)
<ipython-input-8-63fef201e93e> in <module>()
      1 buffer = io.BytesIO()
      2 torch.save(example_model.state_dict(), buffer)
----> 3 loaded = torch.load(buffer)

1 frames
/usr/local/lib/python3.7/dist-packages/torch/serialization.py in _legacy_load(f, map_location, pickle_module, **pickle_load_args)
    775             "functionality.")
    776 
--> 777     magic_number = pickle_module.load(f, **pickle_load_args)
    778     if magic_number != MAGIC_NUMBER:
    779         raise RuntimeError("Invalid magic number; corrupt file?")

EOFError: Ran out of input

Expected behavior

Loading an object saved to io.BytesIO should be an identity operation.

Environment

(google colab)

PyTorch version: 1.9.0+cu111
Is debug build: False
CUDA used to build PyTorch: 11.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.5 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: 6.0.0-1ubuntu2 (tags/RELEASE_600/final)
CMake version: version 3.12.0
Libc version: glibc-2.26

Python version: 3.7.12 (default, Sep 10 2021, 00:21:48) [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.4.104+-x86_64-with-Ubuntu-18.04-bionic
Is CUDA available: False
CUDA runtime version: 11.1.105
GPU models and configuration: Could not collect
Nvidia driver version: Could not collect
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.5
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.0.5
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.19.5
[pip3] torch==1.9.0+cu111
[pip3] torchsummary==1.5.1
[pip3] torchtext==0.10.0
[pip3] torchvision==0.10.0+cu111
[conda] Could not collect

cc @ezyang @gchanan @zou3519 @bdhirsh @jbschlosser @mruberry

[bdhirsh]

marking hi-pri because I can reproduce the error using the snippet from the docs (technically not, because the docs don't include a call to torch.load(), which is where the error occurs):

>>> import torch
>>> import io
>>> x = torch.tensor([0, 1, 2, 3, 4])
>>> buffer = io.BytesIO()
>>> torch.save(x, buffer)
>>> loaded = torch.load(buffer)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/scratch/hirsheybar/pytorch-env/lib/python3.9/site-packages/torch/serialization.py", line 608, in load
    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)
  File "/scratch/hirsheybar/pytorch-env/lib/python3.9/site-packages/torch/serialization.py", line 777, in _legacy_load
    magic_number = pickle_module.load(f, **pickle_load_args)
EOFError: Ran out of input

[romovpa]

This is not a bug of PyTorch, but the natural behavior of io.BytesIO(). After torch.save writes to the buffer, the stream position is set to its end, hence there is nothing to read.

If you set manually the stream position to the beginning (by using seek(pos) method), it will work as expected.

import io
import torch
import torch.nn as nn

example_model = nn.Sequential(
    nn.Linear(10, 10),
    nn.ReLU(),
    nn.Linear(10, 10),
)

buffer = io.BytesIO()
torch.save(example_model.state_dict(), buffer)

# Set the stream position to the beginning
buffer.seek(0)

loaded = torch.load(buffer)

# loaded = OrderedDict([('0.weight', ...

[bdhirsh]

ah, thanks for clarifying @romovpa. Since this doesn't seem like an actual bug I'll close the issue