wittawasw · December 12, 2022 06:43
diff --git a/set_max_split_size_mb.py b/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")
diff --git a/set_max_split_size_mb.sh b/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
	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")
	# 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