Creating and training a U-Net model with PyTorch for 2D & 3D semantic segmentation: Training [3/4]

A guide to semantic segmentation with PyTorch and the U-Net

Image by author

In the previous chapters we created our dataset and built the U-Net model. Now it is time to start training. For that we will write our own training loop within a simple Trainer class and save it in…