Hi @ptrblck, sorry  need your help again, I am trying to learn how to implement custom class for resnet50 > models.
I am trying this but is giving me an error.
RuntimeError                              Traceback (most recent call last)
<ipython-input-31-44bf8238479a> in <module>
     13     checkpoint = torch.load("../models/resnext50_32x4d-7cdf4587.pth")
     14 
---> 15     model_ft = MyResNeXt().to(device) # models.resnet18(pretrained=True)
     16 
     17     del checkpoint
 RuntimeError: CUDA error: device-side assert triggered

here is the code
import torchvision.models as models

class MyResNeXt(models.resnet.ResNet):
    def __init__(self, training=True):
        super(MyResNeXt, self).__init__(block=models.resnet.Bottleneck,
                                        layers=[3, 4, 6, 3], 
                                        groups=32, 
                                        width_per_group=4)

        self.load_state_dict(checkpoint)

        # Override the existing FC layer with a new one.
        self.fc = nn.Linear(2048, 1)

def freeze_until(net, param_name):
    found_name = False
    for name, params in net.named_parameters():
        if name == param_name:
            found_name = True
        params.requires_grad = found_name

######################################################################
# Finetuning the convnet
# ----------------------
#
# Load a pretrained model and reset final fully connected layer.
#
# loop over the number of models to train
for i in np.arange(0, 10):
    
    # initialize the optimizer and model
    print("[INFO] training model {}/{}".format(i + 1, 10))
    
    checkpoint = torch.load("../models/resnext50_32x4d-7cdf4587.pth")
    
    model_ft = MyResNeXt().to(device) # models.resnet18(pretrained=True)
    
    del checkpoint
    
    freeze_until(model_ft, "layer4.0.conv1.weight")

    criterion = nn.CrossEntropyLoss()

    # Observe that all parameters are being optimized
    optimizer_ft = torch.optim.Adam(model_ft.parameters(), lr=1e-5)

    # Decay LR by a factor of 0.1 every 7 epochs
    exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1)
    
    # Train and evaluate
    model_ft = train_model(model_ft, criterion, optimizer_ft, exp_lr_scheduler, num_epochs=25)
    
    # save the model to disk
    p = ["models", "model_resnet18_adam{}.model".format(i)]
    checkpoint = { "optimizer": optimizer_ft.state_dict(),"model": model_ft.state_dict() }
    torch.save(checkpoint, p)

Thanks again for your help!!! I really have been tryingto find advance info in the web about pytorch, but nothing reall that goes beyond MNIST…