Caching function that returns PyTorch tensor on GPU gives error message

Summary

The function described below causes Streamlit to give an error message that says “In this specific case, it’s very likely you found a Streamlit bug so please file a bug report here.”

Steps to reproduce

On a machine with a CUDA GPU, run the following file test.py with streamlit run test.py:

import streamlit as st
import torch
import numpy as np

@st.cache
def f(w):
    return w

w = torch.from_numpy(np.array([1,2,3])).to('cuda')

st.write(f(w))

If you change cuda to cpu it works fine without an error.

If you change the initialization to w = torch.eye(3).to('cuda') then the error message is slightly different.

Expected behavior:

It should write the array [1,2,3].

Actual behavior:

Streamlit gives an error message and says to report a bug.

Is this a regression?

No, this code caused version 0.57 to core dump.

Debug info

• Streamlit version: 0.60
• Python version: 3.6.9
• Using Conda? PipEnv? PyEnv? Pex? No
• OS version: Linux 5.0.0-37-generic #40~18.04.1-Ubuntu
• Browser version: Chrome

[hertzmann]

[randyzwitch]

Closing, as it appears there is a solution provided above. Happy to re-open if I’m misunderstanding