Skip to content

Instantly share code, notes, and snippets.

@wittawasw

wittawasw/set_max_split_size_mb.py

Last active December 12, 2022 06:43
Show Gist options
  • Save wittawasw/c8e6146a135f494834f23d3eca5bf1ac to your computer and use it in GitHub Desktop.
Save wittawasw/c8e6146a135f494834f23d3eca5bf1ac to your computer and use it in GitHub Desktop.
Download ZIP
pytorch setting: max_split_size_mb
Raw
set_max_split_size_mb.py
import torch
# Use when setting up CUDA in the beginning.
# Set max split size = 256 MB
torch.cuda.memory._set_allocator_settings("max_split_size_mb:256")
Raw
set_max_split_size_mb.sh
# In terminal active session, use export or set var directly in profile
# Set max split size = 256 MB
export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:256
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment