Pytorch embedding too big for GPU but fits in CPU - pytorch-lightning

Question

I am using PyTorch lightning, so lightning control GPU/CPU assignments and in
return I get easy multi GPU support for training.
I would like to create an embedding that does not fit in the GPU memory.
fit_in_cpu = torch.nn.Embedding(too_big_for_GPU, embedding_dim)
Then when I select the subset for a batch, send it to the GPU
GPU_tensor = embedding(idx)
How do I do this in Pytorch Lightning?

Answer 1

Lightning will send anything that is registered as a model parameter to GPU, i.e: weights of layers (anything in torch.nn.*) and variables registered using torch.nn.parameter.Parameter.
However if you want to declare something in CPU and then on runtime move it to GPU you can go 2 ways:
Create the too_big_for_GPU inside the __init__ without registering it as a model parameter (using torch.zeros or torch.randn or any other init function). Then move it to the GPU on the forward pass
class MyModule(pl.LightningModule):
def __init__():
self.too_big_for_GPU = torch.zeros(4, 1000, 1000, 1000)
def forward(self, x):
# Move tensor to same GPU as x and operate with it
y = self.too_big_for_GPU.to(x.device) * x**2
return y
Create the too_big_for_GPU which will be created by default in CPU and then you would need to move it to GPU
class MyModule(pl.LightningModule):
def forward(self, x):
# Create the tensor on the fly and move it to x GPU
too_big_for_GPU = torch.zeros(4, 1000, 1000, 1000).to(x.device)
# Operate with it
y = too_big_for_GPU * x**2
return y