Yes, here is the colab notebook that I used. and this 2 is source of my model (which automatically downloads the weights). For last layer, I did the following:
model.last_linear = nn.Sequential(
    nn.Linear(1536, 700),
    nn.ReLU(inplace=True),
    nn.Linear(700, 35)
)