clear_mem.py

#!/usr/bin/env python3
import torch
import gc

shape = (10,000)

input = torch.ones((shape, shape), device="cuda")


def clear_memory(model):
    model.to('cpu')
    gc.collect()
    torch.cuda.empty_cache()
    torch.cuda.ipc_collect()
    torch.clear_autocast_cache()

for _ in range(6):
    linear = torch.nn.Linear(shape, shape).to("cuda")
    output = linear(input)

    clear_memory(linear)