I need some help. Basically I want to get the normal output of the ResNet18 that I modified, but also the one from the nn.Linear(512,32) I am able to get the one before nn.Linear(512,32), but not the one coming out, so the one just before the nn.Linear(32,5). I need this to then use an LSTM with these extracted features.
import torch
from torchvision import models
import torch.nn as nn


class Resnet18(torch.nn.Module):
    def __init__(self):
        super(Resnet18, self).__init__()
        resnet18_pretrained = models.resnet18(pretrained=True)

        self.model = resnet18_pretrained
        self.model.fc =  nn.Sequential(nn.Linear(512, 32),nn.Linear(32, 5))
        self.couches_before_fc = list(self.model.children())[:-1]
        self.resnet_before_fc = nn.Sequential(*self.couches_before_fc)
        self.resnet_before_fc.fc = self.model.fc[0]
        #self.resnet_before_fc = nn.Sequential(*self.couches_before_fc,self.model.fc[0]) #also try this way

    def forward(self, x):

        before_last_fc = self.resnet_before_fc(x)

        x=self.model(x)

        return x, before_last_fc

This strategy seems logic in my mind but I got the following error:
RuntimeError: CUDA error: CUBLAS_STATUS_INVALID_VALUE when calling cublasSgemm( handle, opa, opb, m, n, k, &alpha, a, lda, b, ldb, &beta, c, ldc) related to the line before_last_fc = self.resnet_before_fc(x).
By using a hook I am able to get the value of the output before the nn.Linear(32, 5) but I am not able to get it in a tensor form and return it at the end of the forward…
self.model.fc[0].register_forward_hook(lambda m, input, output: print(output))

also this additional line make the code crashing after trying to save the model:
torch.save(model, model_out_path)
AttributeError: Can’t pickle local object ‘Resnet18.forward..’
Thanks a lot!