Obtain a set of embedding from pretrained model - vgg16 pytorch
For a certain project purpose I am trying to store the 1 * 4096 embeddings (The output right before the final layer) of around 6000 images into a pkl file. For the same, I am running an iteration over the 6000 images on vgg16 modified model in google colab. But it returns 'CUDA out of memory. Tried to allocate 14.00 MiB (GPU 0; 15.90 GiB total capacity; 14.86 GiB already allocated; 1.88 MiB free; 342.26 MiB cached)' error. 

Whereas I have used the same dataset split into test-train for training and validating my model and that runs fine. I am wondering why obtaining and storing the embedding alone is becoming a heavy task in colab. 
Is there any other way I can obtain the embeddings and store in a pkl file other than the below code.

embedding = []
vgg16 = vgg16.to(device)
for x in range (0, len(inputImages)) :
  input = transformations(inputImages[x]) //pre processing
  input = torch.unsqueeze(input, 0)
  input = input.to(device)
  embedding.append(vgg16(input))


The code is interupted at the last line with the CUDA out of memory error. 

The output that you have generated vgg16(input), thats still in cuda. This is so because this output is used for calculating the loss afterwards. So to avoid having your output being stored in CUDA and eat up your GPU memory, move it to CPU using .cpu().numpy(). If that throws an error, you might have to use .detach() as well to detach the variable.
