I’m trying to use LSTM on image data and quickly running out of device memory. My loop runs something like this:
x = torch.rand(20, 2, 3, 224, 738)  # T x B x CHW
# init Variables h and c to proper size
y = []
for xt in x:
    xt = xt.cuda()
    yt, (h,c) = lstm(xt, (h, c))
    y.append(yt)

What I want to achieve is this: take a time slice from x (called xt above), transfer it to device, perform the LSTM operation on GPU, and then get the xt back to host, preserving the Variable connections.
Any ideas?
Thanks