Thanks for the crazy fast reply 
Sure, here is a gist with some code (sorry for the many lines):

  
      gist.github.com
  
  
    https://gist.github.com/PrimeF/2f49e6ffb5f0be279250907aca886220 5
dataset.py
# Python standard libraries
import os
import glob

# Installed libraries
import numpy as np
from PIL import Image, ImageCms
import torch
from torch.utils.data import DataLoader
from torch.utils.data.dataset import Dataset
This file has been truncated. show original
model.py
# Installed libraries
import torch
import torch.nn as nn
from torch.cuda.amp import autocast

# Custom libraries
from utils import resolve_class


class Model(nn.Module):
This file has been truncated. show original
train.py
def train(data_loader_, model_, optimizer, optim_params, loss_fn, epochs, checkpoint_freq, model_dir_, checkpoint=None):
    """
    Trains the model

    :param data_loader_: DataLoader to retrieve the training data
    :param model_: Model to train
    :param optimizer: Optimizer
    :param loss_fn: Loss function
    :param lr: Learning rate
    :param epochs: Number of epochs to train
This file has been truncated. show original




  
  
    
    
  
  


I can post the entire code if you want, I just thought it might be too much code.