Pytorch embedding too big for GPU but fits in CPU
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?

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

