Retraining pytorch model (augmented learning)

I have the following code:

import torch
from facenet_pytorch import InceptionResnetV1, MTCNN
from import DataLoader
from torchvision import datasets
import numpy as np
import pandas as pd
import os

workers = 0 if == 'nt' else 4
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print('Running on device: {}'.format(device))

mtcnn = MTCNN(
    image_size=160, margin=0, min_face_size=20,
    thresholds=[0.6, 0.7, 0.7], factor=0.709, post_process=True,

def collate_fn(x):
    return x[0]

dataset = datasets.ImageFolder('data/images/')
dataset.idx_to_class = {i:c for c, i in dataset.class_to_idx.items()}
loader = DataLoader(dataset, collate_fn=collate_fn, num_workers=workers)

aligned = []
names = []
i = 0
for x, y in loader:
    x_aligned, prob = mtcnn(x, return_prob=True)
    if x_aligned is not None:
        print('Face detected with probability: {:8f}'.format(prob))
        i += 1

for name, param in mtcnn.named_parameters(): #Freezing everything but last layer
    if name != "onet.dense6_3.bias":
        param.require_grad = False
        param.require_grad = True

And now I would like to retrain this model to predict three classes (Now it only predicts the probability of a face). Let say that I have inside data/images/ three folders, faces1, faces2 and faces3. How could I retrain this model with these three folders? I would like to have a tensor like [prob1, prob2, prob3] with the probability of an image for each class. Thanks.