mem_usage.py

import sklearn
import numpy as np
import torch


sklearn.set_config(array_api_dispatch=True)

def my_code(X, cdist=False):
    if cdist:
        dist = torch.cdist(X, X, p=2)
    else:
        dist = sklearn.metrics.pairwise.euclidean_distances(X, X)
    print(dist.shape)
    return dist


torch.cuda.memory._record_memory_history(
    max_entries=100_000
)

X = np.random.randn(50_000, 36)
X_t = torch.asarray(X, device="cuda")

cdist = True

for n in range(5):
    my_code(X_t, cdist=cdist)

if cdist:
    torch.cuda.memory._dump_snapshot(f"mem_usage_cdist.pickle")
else:
    torch.cuda.memory._dump_snapshot(f"mem_usage.pickle")
torch.cuda.memory._record_memory_history(enabled=None)