Solved by RaLo4 in post #2 
              
              
                This might not be the most elegant way but it definitely works 
if not(os.path.isdir('/ImageOutput')):
        os.mkdir('/ImageOutput')
torchvision.utils.save_image(code, f"/ImageOutput/code{epoch}.png", nrow=num_rows)