If anyone comes across this issue. I found out that refactoring the loss code into its own function resolves the problem.
Probably TensorFlow optimizes the graph structure, while the previous PyTorch-code still contains a reference to the resized image so PyTorch keeps it in memory, this is not the case after the refactoring.
After the change, the loss looks more or less the same.

Selection_0352103×1201 51.9 KB

Here is the fixed code:
import torch
import torch.nn.functional as F
import sys
from math import log2
import numpy as np
from PIL import Image

_suffixes = ["bytes", "KiB", "MiB", "GiB", "TiB", "PiB", "EiB", "ZiB", "YiB"]


def file_size(size):
    order = int(log2(size) / 10) if size else 0
    return "{:.4g} {}".format(size / (1 << (order * 10)), _suffixes[order])


def get_loss(image, resizescale=4):
    # print(f"current mem: {file_size(torch.cuda.memory_allocated(0))}")
    w = 256
    h = 256
    print(f"Resizing with resize scale: {resizescale}")
    resized = F.interpolate(
        image,
        size=[int(resizescale * h), int(resizescale * w)],
        mode="bilinear",
        align_corners=True,
    )

    print(f"current size: {resized.shape}")

    sum = resized.sum()
    return sum


def main(resizescale=4, save=True):
    # create an image tensor
    device = torch.device("cuda")

    print(f"Creating image")
    np.random.seed(42)
    # np_image = np.random.rand(1, 1, h, w)
    np_image = np.array(Image.open("test.png")).transpose((2, 0, 1))[None, ...]
    print(f"geting sum")

    image = torch.tensor(np_image, device=device, dtype=torch.float32)
    image.requires_grad = True
    sum = get_loss(image, resizescale=resizescale)
    sum.backward()
    # if save:
    #     np.save("torch_image.npy", resized.detach().cpu().numpy())
    #     np.save("torch_grad.npy", image.grad.detach().cpu().numpy())

    print(f"current mem: {file_size(torch.cuda.max_memory_allocated(0))}")
    return torch.cuda.max_memory_allocated(0)