Yes sure.
The transforamations I used are these ones:
data_transforms = {

    'train': transforms.Compose([

        transforms.RandomResizedCrop(224),

        transforms.RandomHorizontalFlip(),

        transforms.ToTensor()]),

    'val': transforms.Compose([

        transforms.Resize(256),

        transforms.CenterCrop(224),

        transforms.ToTensor()])

}

The functions:
def train_model(model, dataloaders, criterion, optimizer, scheduler, batch_size=5, num_epochs=10):#

    since = time.time()

    val_acc_history = []

    best_model_wts = copy.deepcopy(model.state_dict())

    best_acc = 0.0

    #pdb.set_trace()

    for epoch in range(num_epochs):

        print('Epoch {}/{}'.format(epoch, num_epochs - 1))

        print('-' * 10)

        # Each epoch has a training and validation phase

        for phase in ['train']:#, 'val']:

            if phase == 'train':

                model.train()  # Set model to training mode

            running_loss = 0.0

            running_corrects = 0

            average_precis_train = 0.001

            average_precis_train_per_class = 0.001

            loss_values = []

            gr_truth_array = np.array([]) #convet to int dtype

            preds_array = np.array([])

            gr_truth_array = gr_truth_array.astype(int)

            preds_array = preds_array.astype(int) 

            average_precision_array = np.array([]).astype(float)


            print('Iterating over data:')

            for batch_idx, (inputs, labels) in enumerate(dataloaders[phase]):

                inputs = inputs.to(device)

                labels = labels.to(device).float()

                gt_data = labels

                gt_data = gt_data.to(device)

                gt_data = gt_data.cpu().data.numpy()

                #average_precision_array = []

                # zero the parameter gradients

                optimizer.zero_grad()

                # forward

                # track history if only in train

                #pdb.set_trace()

                if phase == 'train':

                  with torch.set_grad_enabled(phase == 'train'):

                      outputs = model(inputs)

                      outputs = outputs.cpu()#.data.numpy()

                      preds = outputs.cpu().data.numpy()

                      preds = np.round(preds) #set a condition for binary

                      preds_int = preds.astype(int)

                      gt_data_np = np.round(gt_data)

                      gt_data_int = gt_data_np.astype(int)

                      gt_data = torch.from_numpy(gt_data_np)

                      loss = criterion(outputs, gt_data)

                      gr_truth_array = np.append(gr_truth_array, gt_data_int)

                      preds_array = np.append(preds_array ,preds_int)

                    # backward + optimize only if in training phase

                      if phase == 'train':

                          loss.backward()

                          optimizer.step()

                # statistics

                  gr_truth_array = np.reshape(gr_truth_array, (-1, 40))

                  preds_array = np.reshape(preds_array, (-1, 40))

                  running_loss += loss.item() * inputs.size(0)

                  running_corrects += f1_score(gt_data, preds, average="samples")

                

            if phase == 'train':

                scheduler.step()

                average_precis_train += average_precision_score(gr_truth_array, preds_array, average= "macro")

                average_precis_train_per_class += average_precision_score(gr_truth_array, preds_array, average=None)

                average_precision_array = np.append(average_precision_array, average_precis_train_per_class)

                #pdb.set_trace()

                av_precis_array = [j for i in zip(average_precision_array, attributes) for j in i]

                av_precis_array = np.array(av_precis_array)

                print("Average precision Training:", average_precis_train)      

                print("Average precision per Class Training:", av_precis_array)

            #pdb.set_trace()   

            epoch_loss = running_loss / len(dataloaders[phase].dataset)

            epoch_acc = running_corrects / len(dataloaders[phase].dataset) #running_corrects.float()

            epoch_acc = np.round(epoch_acc, decimals=4)

             

            print('{} Loss: {:.4f}'.format(phase, epoch_loss))

            print("Acc:", epoch_acc)

 

    time_elapsed = time.time() - since

    print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))

    model.load_state_dict(best_model_wts)

    return model, val_acc_history

The evaluation is almost the same but the model is set to model.eval and I use with torch.no_grad(): instead of set_grad_enabled