Thanks for the reply.
Upon rerunning the code from op the code does not crash for unknown reasons,
but when we apply autocast it produces unexpected results.
This is the essence of the training and evaluation code.
scaler = torch.cuda.amp.GradScaler()
for epoch in range(epochs):
    for images in train_photo_dataloader:
        images = images.to(device)

        optimizer.zero_grad()

        with torch.cuda.amp.autocast():
            fakes = G(images)
            loss = content_loss(fakes, images)
        scaler.scale(loss).backward()

        scaler.step(optimizer)

        scaler.update()

    G.eval()
    with torch.no_grad():
        fakes = G(test_images)
        loss_test = content_loss(fakes, test_images)

        logger.log_images("Test images output", fakes, normalize=False)
    G.train()

With the logger using make_grid like this
class ExperimentLogger(object):
    def __init__(self, exp_id: str, init_phase: bool):
        self.global_step = 0
        self.exp_id = exp_id
        log_dir = f"./logs/{exp_id}/{'pre_training' if init_phase else 'main'}"
        self.writer = SummaryWriter(log_dir=log_dir)

    def log_images(self, tag: str, tensor: torch.Tensor, normalize=False):
        tensor = vutils.make_grid(tensor.cpu(), normalize=normalize)
        # TODO Maybe check here what phase we are in and set 'tag' accordingly.
        self.writer.add_image(tag, tensor, self.global_step)

G is a Unet defined using segmentation_models.pytorch
smp.Unet(encoder_name="resnet18", encoder_depth=5, encoder_weights="imagenet",
                         decoder_channels=(32, 32, 32, 32, 16),
                         decoder_use_batchnorm=True, in_channels=3, classes=3, activation="sigmoid")

When using autocast this is our result individualImage (1)906×228 194 KB
compared to not not wrapping the two code lines in with torch.cuda.amp.autocast():
individualImage (2)906×228 423 KB
We are simply trying to recreate the following images individualImage906×228 128 KB
Another problem is that the produced are oversaturated, compared to the originals.