Thanks for the help!
where should I put that? Becasue it gives me this error:
 File "trainer.py", line 322, in main
    for ii,module in enumerate(self.resnet):
NameError: name 'self' is not defined

I have my code divided into two parts: resnet and trainer:
Trainer.py:
import argparse
import os
import shutil
import time

import pickle
import numpy as np
import pandas as pd

import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim
import torch.utils.data
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import resnet

from torch.utils.data import Dataset
from torchvision import transforms
from PIL import Image
from imgaug import augmenters as iaa
from matplotlib import pyplot as plt
import torchvision.transforms.functional as TF
import yaml

model_names = sorted(name for name in resnet.__dict__
    if name.islower() and not name.startswith("__")
                     and name.startswith("resnet")
                     and callable(resnet.__dict__[name]))

print(model_names)

parser = argparse.ArgumentParser(description='Propert ResNets for CIFAR10 in pytorch')
parser.add_argument('--arch', '-a', metavar='ARCH', default='resnet32',
                    choices=model_names,
                    help='model architecture: ' + ' | '.join(model_names) +
                    ' (default: resnet32)')
parser.add_argument('-j', '--workers', default=4, type=int, metavar='N',
                    help='number of data loading workers (default: 4)')
parser.add_argument('--epochs', default=200, type=int, metavar='N',
                    help='number of total epochs to run')
parser.add_argument('--start-epoch', default=0, type=int, metavar='N',
                    help='manual epoch number (useful on restarts)')
parser.add_argument('-b', '--batch-size', default=128, type=int,
                    metavar='N', help='mini-batch size (default: 128)')
parser.add_argument('--lr', '--learning-rate', default=0.1, type=float,
                    metavar='LR', help='initial learning rate')
parser.add_argument('--momentum', default=0.9, type=float, metavar='M',
                    help='momentum')
parser.add_argument('--weight-decay', '--wd', default=1e-4, type=float,
                    metavar='W', help='weight decay (default: 5e-4)')
parser.add_argument('--print-freq', '-p', default=50, type=int,
                    metavar='N', help='print frequency (default: 20)')
parser.add_argument('--resume', default='', type=str, metavar='PATH',
                    help='path to latest checkpoint (default: none)')
parser.add_argument('-e', '--evaluate', dest='evaluate', action='store_true',
                    help='evaluate model on validation set')
parser.add_argument('--pretrained', dest='pretrained', action='store_true',
                    help='use pre-trained model')
parser.add_argument('--half', dest='half', action='store_true',
                    help='use half-precision(16-bit) ')
parser.add_argument('--save-dir', dest='save_dir',
                    help='The directory used to save the trained models',
                    default='save_temp', type=str)
parser.add_argument('--save-every', dest='save_every',
                    help='Saves checkpoints at every specified number of epochs',
                    type=int, default=10)
best_prec1 = 0

CONFIG = open("/content/drive/My Drive/CIFAR10_pruebas/pytorch_resnet_cifar10/config_CIFAR10.yaml")
CONFIG2 = open("/content/drive/My Drive/CIFAR10_pruebas/pytorch_resnet_cifar10/config_ResNet32.yaml")
CONFIG = yaml.load(CONFIG)
CONFIG2 = yaml.load(CONFIG2)

res_path = "/content/drive/My Drive/CIFAR10_pruebas/pytorch_resnet_cifar10/save_graficas/4"

class CifarDataset(Dataset):
    """Class for ADE20K dataset."""

    def __init__(self, root_dir, train, CONFIG, imbalance_factor=1, ten_Crop=False):
        """
        Initialize the dataset
        :param root_dir: Root directory to the dataset
        :param set: Dataset set: Training or Validation
        :param clean: Use the cleaned version of ADE20K instead of the original one
        """
        # Extract main path and set (Train or Val)
        self.image_dir = os.path.join(root_dir, CONFIG['DATASET']['NAME'])
        self.train = train
        self.prog_sprinkles = CONFIG['DATASET']['PROG_SPRINKLES']

        # Decode dataset scene categories
        self.classes = list()
        class_file_name = os.path.join(self.image_dir, "classes.txt")

        with open(class_file_name) as class_file:
            for line in class_file:
                self.classes.append(line.split()[0])

        self.nclasses = self.classes.__len__()

        # Create list for filenames and scene ground-truth labels
        self.filenames = list()
        self.labels = list()
        self.labelsindex = list()

        if self.train:
            filenames_file = os.path.join(self.image_dir, ("train2.txt"))
        else:
            filenames_file = os.path.join(self.image_dir, "val.txt")

        with open(filenames_file) as class_file:
            for line in class_file:
                label, name = line.split('/')
                self.filenames.append(name.split()[0])
                self.labels.append(label)
                self.labelsindex.append(self.classes.index(label))

        # Control Statements for data loading
        assert len(self.filenames) == len(self.labels)

        # ----------------------------- #
        #     ImAug Transformations     #
        # ----------------------------- #
        # Transformations for train set
        self.ImAugTraforms = iaa.Sequential([
        #     # Small gaussian blur with random sigma between 0 and 0.5.
            iaa.Sometimes(0.5, iaa.GaussianBlur(sigma=(0, 0.5))),
        #     # Strengthen or weaken the contrast in each image.
            iaa.LinearContrast((0.7, 1.5)),
        #     # Add gaussian noise.
            iaa.AdditiveGaussianNoise(loc=0, scale=(0.0, 0.05 * 255), per_channel=0.5),
        #     # Make some images brighter and some darker.
            iaa.Multiply((0.75, 1.3), per_channel=0.2),
            iaa.Cutout(nb_iterations=1, size=0.05, fill_mode="constant", cval=0)
        #     # Rotate iamge between -25 and 25
        #     iaa.Affine(rotate=(-25, 25)),
        ], random_order=True)  # apply augmenters in random order
        #self.ImAugTraforms = iaa.Cutout(nb_iterations=1, size=0.05, fill_mode="constant", cval=0)

        # ----------------------------- #
        #    Pytorch Transformations    #
        # ----------------------------- #
        # Define Random crop. If image is smaller resize first.
        if self.train:
            self.transforms = transforms.Compose(
                [transforms.Pad(CONFIG2['MODEL']['PADDING']),
                 transforms.RandomCrop(CONFIG2['MODEL']['CROP']),
                 transforms.RandomHorizontalFlip(),
                 transforms.ToTensor(),
                 transforms.Normalize(CONFIG['DATASET']['MEAN'], CONFIG['DATASET']['STD'])])
        else:
            if ten_Crop:
                self.transforms = transforms.Compose(
                    [transforms.TenCrop(CONFIG2['MODEL']['CROP']),
                     transforms.Lambda(lambda crops: torch.stack([transforms.ToTensor()(crop) for crop in crops])),
                     transforms.Lambda(lambda crops: torch.stack([transforms.Normalize(CONFIG['DATASET']['MEAN'],
                                                                                       CONFIG['DATASET']['STD'])
                                                                  (crop) for crop in crops]))])
            else:
                self.transforms = transforms.Compose(
                    [transforms.CenterCrop(CONFIG2['MODEL']['CROP']),
                     transforms.ToTensor(),
                     transforms.Normalize(CONFIG['DATASET']['MEAN'], CONFIG['DATASET']['STD'])])

    def __len__(self):
        """
        Function to get the size of the dataset
        :return: Size of dataset
        """
        return len(self.filenames)

    def __getitem__(self, idx):
        """
        Function to get a sample from the dataset. First both RGB and Semantic images are read in PIL format. Then
        transformations are applied from PIL to Numpy arrays to Tensors.

        For regular usage:
            - Images should be outputed with dimensions (3, W, H)
            - Semantic Images should be outputed with dimensions (1, W, H)

        In the case that 10-crops are used:
            - Images should be outputed with dimensions (10, 3, W, H)
            - Semantic Images should be outputed with dimensions (10, 1, W, H)

        :param idx: Index
        :return: Dictionary containing {RGB image, semantic segmentation mask, scene category index}
        """

        # Get RGB image path and load it
        if self.train:
            img_name = os.path.join(self.image_dir, "train", self.labels[idx], self.filenames[idx])
        else:
            img_name = os.path.join(self.image_dir, "val", self.labels[idx], self.filenames[idx])

        img = Image.open(img_name)

        # ImAug Transforms
        if self.train and self.prog_sprinkles:
            img = Image.fromarray(self.ImAugTraforms.augment_image(np.asarray(img)))

        # TAREA PRETEXTO -- ROTACIÓN
        ## List of 4 images, each with a different rotation
        #rotated_imgs = [
        #    self.transforms(img), ## 0 degrees rotation
        #    self.transforms(Image.fromarray(rotate_img(img,90).copy())), ## 90 degrees rotation
        #    self.transforms(Image.fromarray(rotate_img(img,180).copy())), ## 180 degrees rotation
        #    self.transforms(Image.fromarray(rotate_img(img,270).copy())) ## 270 degrees rotation
        #]

        ## Createe rotation labels, i.e. 4 (0-3)
        #rot_labels = torch.LongTensor([0, 1, 2, 3])

        ## Stack the list of rotated images into a tensor
        #rot_img = torch.stack(rotated_imgs, dim=0)

        # PyTorch Transforms
        img = self.transforms(img)
        # Create dictionary
        #self.sample = {'Index': idx, 'Image': img, 'Class': self.labelsindex[idx]}

        return img, self.labelsindex[idx]

    def redefineSprinkles(self, epoch):
        """
        Function to apply curriculum learning to Progressive Sprinkles. Depending on the epoch sprinkles will be
        bigger and higher in number making classification more difficult.
        """
        if epoch < 20:
            print("Época:",epoch)
            return iaa.Sequential([
                iaa.Sometimes(0.5, iaa.GaussianBlur(sigma=(0, 0.5))),
                iaa.LinearContrast((0.7, 1.5)),
                iaa.AdditiveGaussianNoise(loc=0, scale=(0.0, 0.05 * 255), per_channel=0.5),
                iaa.Multiply((0.75, 1.3), per_channel=0.2),
                iaa.Cutout(nb_iterations=1, size=0.05, fill_mode="constant", cval=0)
            ], random_order=True)
        elif epoch < 40:
            print("Época:",epoch)
            return iaa.Sequential([
                iaa.Sometimes(0.5, iaa.GaussianBlur(sigma=(0, 0.5))),
                iaa.LinearContrast((0.7, 1.5)),
                iaa.AdditiveGaussianNoise(loc=0, scale=(0.0, 0.05 * 255), per_channel=0.5),
                iaa.Multiply((0.75, 1.3), per_channel=0.2),
                iaa.Cutout(nb_iterations=2, size=0.05, fill_mode="constant", cval=0)
            ], random_order=True)
        elif epoch < 60:
            print("Época:",epoch)
            return iaa.Sequential([
                iaa.Sometimes(0.5, iaa.GaussianBlur(sigma=(0, 0.5))),
                iaa.LinearContrast((0.7, 1.5)),
                iaa.AdditiveGaussianNoise(loc=0, scale=(0.0, 0.05 * 255), per_channel=0.5),
                iaa.Multiply((0.75, 1.3), per_channel=0.2),
                iaa.Cutout(nb_iterations=3, size=0.1, fill_mode="constant", cval=0)
            ], random_order=True)
        elif epoch < 80:
            print("Época:",epoch)
            return iaa.Sequential([
                iaa.Sometimes(0.5, iaa.GaussianBlur(sigma=(0, 0.5))),
                iaa.LinearContrast((0.7, 1.5)),
                iaa.AdditiveGaussianNoise(loc=0, scale=(0.0, 0.05 * 255), per_channel=0.5),
                iaa.Multiply((0.75, 1.3), per_channel=0.2),
                iaa.Cutout(nb_iterations=4, size=0.1, fill_mode="constant", cval=0)
            ], random_order=True)
        elif epoch < 100:
            print("Época:",epoch)
            return iaa.Sequential([
                iaa.Sometimes(0.5, iaa.GaussianBlur(sigma=(0, 0.5))),
                iaa.LinearContrast((0.7, 1.5)),
                iaa.AdditiveGaussianNoise(loc=0, scale=(0.0, 0.05 * 255), per_channel=0.5),
                iaa.Multiply((0.75, 1.3), per_channel=0.2),
                iaa.Cutout(nb_iterations=5, size=0.2, fill_mode="constant", cval=0)
            ], random_order=True)
        else:
            print("Época:",epoch)
            return iaa.Sequential([
                iaa.Sometimes(0.5, iaa.GaussianBlur(sigma=(0, 0.5))),
                iaa.LinearContrast((0.7, 1.5)),
                iaa.AdditiveGaussianNoise(loc=0, scale=(0.0, 0.05 * 255), per_channel=0.5),
                iaa.Multiply((0.75, 1.3), per_channel=0.2),
                iaa.Cutout(nb_iterations=6, size=0.25, fill_mode="constant", cval=0)
            ], random_order=True)


def main():
    global args, best_prec1
    args = parser.parse_args()
    loss_train_epoch = []
    loss_test_epoch = []
    acc_train_per_epoch = []
    acc_test_per_epoch = []


    # Check the save_dir exists or not
    if not os.path.exists(args.save_dir):
        os.makedirs(args.save_dir)

    model = torch.nn.DataParallel(resnet.__dict__[args.arch]())

    checkpoint_path = "/content/drive/My Drive/CIFAR10_pruebas/pytorch_resnet_cifar10/save_resnet32/checkpoint3_1.th"

    # optionally resume from a checkpoint
    if checkpoint_path:
        if os.path.isfile(checkpoint_path):
            print("=> loading checkpoint '{}'".format(checkpoint_path))
            checkpoint = torch.load(checkpoint_path)
            model.load_state_dict(checkpoint['state_dict'])
            print("=> loaded checkpoint")
        else:
            print("=> no checkpoint found at '{}'".format(args.resume))

    #Freeze layers
    for param in model.parameters():
        param.requires_grad = False
    
    #Create you own layers to attach
    lin1 = nn.Linear(64, 10)

    #Check the layers of the pretrained model
    for ii,module in enumerate(self.resnet):
            print("module " + str(ii) + " " + str(module))

    model.cuda()

    cudnn.benchmark = True

    train_set = CifarDataset('./datasets', train=True, CONFIG=CONFIG, imbalance_factor=CONFIG['DATASET']['IMBALANCE'])
    val_set = CifarDataset('./datasets', train=False, CONFIG=CONFIG)

    train_loader =  torch.utils.data.DataLoader(dataset=train_set,
                                     batch_size=args.batch_size,
                                     shuffle=True,
                                     num_workers=args.workers,
                                     pin_memory=True)
    val_loader =  torch.utils.data.DataLoader(dataset=val_set,
                                     batch_size=args.batch_size,
                                     shuffle=False,
                                     num_workers=args.workers,
                                     pin_memory=True)

    # define loss function (criterion) and optimizer
    criterion = nn.CrossEntropyLoss().cuda()

    if args.half:
        model.half()
        criterion.half()

    optimizer = torch.optim.SGD(model.parameters(), args.lr,
                                momentum=args.momentum,
                                weight_decay=args.weight_decay)

    #lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer,
    #                                                    milestones=[100, 150], last_epoch=args.start_epoch - 1)
    #lr_scheduler.last_epoch = args.start_epoch - 1   #CUANDO CARGUEMOS CHECKPOINT.

    if args.arch in ['resnet1202', 'resnet110']:
        # for resnet1202 original paper uses lr=0.01 for first 400 minibatches for warm-up
        # then switch back. In this setup it will correspond for first epoch.
        for param_group in optimizer.param_groups:
            param_group['lr'] = args.lr*0.1


    if args.evaluate:
        validate(val_loader, model, criterion)
        return

    for epoch in range(args.start_epoch, args.epochs):

        # train for one epoch
        print('current lr {:.5e}'.format(optimizer.param_groups[0]['lr']))
        loss_per_epoch, top1_train_ac = train(train_loader, model, criterion, optimizer, epoch)
        loss_train_epoch += [loss_per_epoch]
        #lr_scheduler.step()

        # evaluate on validation set
        loss_per_epoch_test, prec1 = validate(val_loader, model, criterion)

        loss_test_epoch += [loss_per_epoch_test]
        acc_train_per_epoch += [top1_train_ac]
        acc_test_per_epoch += [prec1]

        # remember best prec@1 and save checkpoint
        is_best = prec1 > best_prec1
        best_prec1 = max(prec1, best_prec1)

        if epoch > 0 and epoch % args.save_every == 0:
            save_checkpoint({
                'epoch': epoch + 1,
                'state_dict': model.state_dict(),
                'best_prec1': best_prec1,
            }, is_best, filename=os.path.join(args.save_dir, 'checkpoint4.th'))

        save_checkpoint({
            'epoch' : epoch,
            'state_dict': model.state_dict(),
            'best_prec1': best_prec1,
        }, is_best, filename=os.path.join(args.save_dir, 'model4.th'))

        # save losses:
        np.save(res_path + '/' + 'LOSS_epoch_train.npy', np.asarray(loss_train_epoch))
        np.save(res_path + '/' + 'LOSS_epoch_val.npy', np.asarray(loss_test_epoch))

        # save accuracies:
        np.save(res_path + '/' + 'accuracy_per_epoch_train.npy', np.asarray(acc_train_per_epoch))
        np.save(res_path + '/' + 'accuracy_per_epoch_val.npy', np.asarray(acc_test_per_epoch))


def train(train_loader, model, criterion, optimizer, epoch):
    """
        Run one train epoch
    """
    batch_time = AverageMeter()
    data_time = AverageMeter()
    losses = AverageMeter()
    top1 = AverageMeter()

    # switch to train mode
    model.train()

    end = time.time()

    train_loss = []

    for i, (input, target) in enumerate(train_loader):

        # measure data loading time
        data_time.update(time.time() - end)

        target = target.cuda()
        input_var = input.cuda()

        #size = input_var.size()
        #input_var = input_var.view(-1, size[2], size[3], size[4])
        #target = target.view(-1)

        target_var = target
        if args.half:
            input_var = input_var.half()

        # compute output
        output = model(input_var)
        loss = criterion(output, target_var)

        train_loss.append(loss.item())
        
        # compute gradient and do SGD step
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        output = output.float()
        loss = loss.float()
        # measure accuracy and record loss
        prec1 = accuracy(output.data, target)[0]
        losses.update(loss.item(), input.size(0))
        top1.update(prec1.item(), input.size(0))

        # measure elapsed time
        batch_time.update(time.time() - end)
        end = time.time()

        if i % args.print_freq == 0:
            print('Epoch: [{0}][{1}/{2}]\t'
                  'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                  'Data {data_time.val:.3f} ({data_time.avg:.3f})\t'
                  'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
                  'Prec@1 {top1.val:.3f} ({top1.avg:.3f})'.format(
                      epoch, i, len(train_loader), batch_time=batch_time,
                      data_time=data_time, loss=losses, top1=top1))

    return sum(train_loss)/len(train_loss), top1.avg

def validate(val_loader, model, criterion):
    """
    Run evaluation
    """
    batch_time = AverageMeter()
    losses = AverageMeter()
    top1 = AverageMeter()

    # switch to evaluate mode
    model.eval()

    end = time.time()

    test_loss = []

    with torch.no_grad():
        for i, (input, target) in enumerate(val_loader):

            target = target.cuda()
            input_var = input.cuda()

            #size = input_var.size()
            #input_var = input_var.view(-1, size[2], size[3], size[4])
            #target = target.view(-1)

            target_var = target
            if args.half:
                input_var = input_var.half()

            # compute output
            output = model(input_var)
            loss = criterion(output, target_var)

            test_loss.append(loss.item())

            output = output.float()
            loss = loss.float()

            # measure accuracy and record loss
            prec1 = accuracy(output.data, target)[0]
            losses.update(loss.item(), input.size(0))
            top1.update(prec1.item(), input.size(0))

            # measure elapsed time
            batch_time.update(time.time() - end)
            end = time.time()

            if i % args.print_freq == 0:
                print('Test: [{0}/{1}]\t'
                      'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                      'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
                      'Prec@1 {top1.val:.3f} ({top1.avg:.3f})'.format(
                          i, len(val_loader), batch_time=batch_time, loss=losses,
                          top1=top1))

    print(' * Prec@1 {top1.avg:.3f}'
          .format(top1=top1))

    return sum(test_loss)/len(test_loss), top1.avg

def save_checkpoint(state, is_best, filename='checkpoint.pth.tar'):
    """
    Save the training model
    """
    torch.save(state, filename)

class AverageMeter(object):
    """Computes and stores the average and current value"""
    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count


def accuracy(output, target, topk=(1,)):
    """Computes the precision@k for the specified values of k"""
    maxk = max(topk)
    batch_size = target.size(0)

    _, pred = output.topk(maxk, 1, True, True)
    pred = pred.t()
    correct = pred.eq(target.view(1, -1).expand_as(pred))

    res = []
    for k in topk:
        correct_k = correct[:k].view(-1).float().sum(0)
        res.append(correct_k.mul_(100.0 / batch_size))
    return res


def rotate_img(img, rot):
    if rot == 0:  # 0 degrees rotation
        return img
    elif rot == 90:  # 90 degrees rotation
        return np.flipud(np.transpose(img, (1, 0, 2)))
    elif rot == 180:  # 90 degrees rotation
        return np.fliplr(np.flipud(img))
    elif rot == 270:  # 270 degrees rotation / or -90
        return np.transpose(np.flipud(img), (1, 0, 2))
    else:
        raise ValueError('rotation should be 0, 90, 180, or 270 degrees')


if __name__ == '__main__':
    main()


Resnet:
'''
Properly implemented ResNet-s for CIFAR10 as described in paper [1].

The implementation and structure of this file is hugely influenced by [2]
which is implemented for ImageNet and doesn't have option A for identity.
Moreover, most of the implementations on the web is copy-paste from
torchvision's resnet and has wrong number of params.

Proper ResNet-s for CIFAR10 (for fair comparision and etc.) has following
number of layers and parameters:

name      | layers | params
ResNet20  |    20  | 0.27M
ResNet32  |    32  | 0.46M
ResNet44  |    44  | 0.66M
ResNet56  |    56  | 0.85M
ResNet110 |   110  |  1.7M
ResNet1202|  1202  | 19.4m

which this implementation indeed has.

Reference:
[1] Kaiming He, Xiangyu Zhang, Shaoqing Ren, Jian Sun
    Deep Residual Learning for Image Recognition. arXiv:1512.03385
[2] https://github.com/pytorch/vision/blob/master/torchvision/models/resnet.py

If you use this implementation in you work, please don't forget to mention the
author, Yerlan Idelbayev.
'''
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.nn.init as init

from torch.autograd import Variable

__all__ = ['ResNet', 'resnet20', 'resnet32', 'resnet44', 'resnet56', 'resnet110', 'resnet1202']

def _weights_init(m):
    classname = m.__class__.__name__
    #print(classname)
    if isinstance(m, nn.Linear) or isinstance(m, nn.Conv2d):
        init.kaiming_normal_(m.weight)

class LambdaLayer(nn.Module):
    def __init__(self, lambd):
        super(LambdaLayer, self).__init__()
        self.lambd = lambd

    def forward(self, x):
        return self.lambd(x)


class BasicBlock(nn.Module):
    expansion = 1

    def __init__(self, in_planes, planes, stride=1, option='A'):
        super(BasicBlock, self).__init__()
        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(planes)
        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(planes)

        self.shortcut = nn.Sequential()
        if stride != 1 or in_planes != planes:
            if option == 'A':
                """
                For CIFAR10 ResNet paper uses option A.
                """
                self.shortcut = LambdaLayer(lambda x:
                                            F.pad(x[:, :, ::2, ::2], (0, 0, 0, 0, planes//4, planes//4), "constant", 0))
            elif option == 'B':
                self.shortcut = nn.Sequential(
                     nn.Conv2d(in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False),
                     nn.BatchNorm2d(self.expansion * planes)
                )

    def forward(self, x):
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        out += self.shortcut(x)
        out = F.relu(out)
        return out


class ResNet(nn.Module):
    def __init__(self, block, num_blocks, num_classes=10):
        super(ResNet, self).__init__()
        self.in_planes = 16

        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(16)
        self.layer1 = self._make_layer(block, 16, num_blocks[0], stride=1)
        self.layer2 = self._make_layer(block, 32, num_blocks[1], stride=2)
        self.layer3 = self._make_layer(block, 64, num_blocks[2], stride=2)
        self.linear = nn.Linear(64, num_classes)

        self.apply(_weights_init)

    def _make_layer(self, block, planes, num_blocks, stride):
        strides = [stride] + [1]*(num_blocks-1)
        layers = []
        for stride in strides:
            layers.append(block(self.in_planes, planes, stride))
            self.in_planes = planes * block.expansion

        return nn.Sequential(*layers)

    def forward(self, x):
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.layer1(out)
        out = self.layer2(out)
        out = self.layer3(out)
        out = F.avg_pool2d(out, out.size()[3])
        out = out.view(out.size(0), -1)
        out = self.linear(out)
        return out


def resnet20():
    return ResNet(BasicBlock, [3, 3, 3])


def resnet32():
    return ResNet(BasicBlock, [5, 5, 5])


def resnet44():
    return ResNet(BasicBlock, [7, 7, 7])


def resnet56():
    return ResNet(BasicBlock, [9, 9, 9])


def resnet110():
    return ResNet(BasicBlock, [18, 18, 18])


def resnet1202():
    return ResNet(BasicBlock, [200, 200, 200])


def test(net):
    import numpy as np
    total_params = 0

    for x in filter(lambda p: p.requires_grad, net.parameters()):
        total_params += np.prod(x.data.numpy().shape)
    print("Total number of params", total_params)
    print("Total layers", len(list(filter(lambda p: p.requires_grad and len(p.data.size())>1, net.parameters()))))


if __name__ == "__main__":
    for net_name in __all__:
        if net_name.startswith('resnet'):
            print(net_name)
            test(globals()[net_name]())
            print()

I’m using resnet32.
Thank you so much