Rgpu: a PyTorch device whose tensors live on a remote GPU

Hi all. I’ve been building rgpu, which lets you write device="rgpu" on a laptop while the tensors and kernels live on a GPU host you reach over SSH. pip install rgpu import torch, rgpu x = torch.ones(4, device="rgpu") print((x * 2).sum().item()) # 8.0, computed on the remote GPU How it works rgpu is a PrivateUse1 backend with a RemoteTensor wrapper subclass. Each op’s output shape is worked out locally on meta tensors, so the op returns at once with a handle, and it’s queued and sent to th...