Skip to content

Instantly share code, notes, and snippets.

Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
@sparkydogX
sparkydogX / occupy-memory.py
Last active March 27, 2024 05:55
Pytorch trick : occupy all GPU memory in advance
import os
import torch
from tqdm import tqdm
import time
# declare which gpu device to use
cuda_device = '0'
def check_mem(cuda_device):
devices_info = os.popen('"/usr/bin/nvidia-smi" --query-gpu=memory.total,memory.used --format=csv,nounits,noheader').read().strip().split("\n")