Memory optimizations. Default to using cudamalloc when torch 2.0 for mem allocation

This commit is contained in:
Jaret Burkett
2023-09-12 04:30:23 -06:00
parent e8583860ad
commit d74dd636ee
5 changed files with 104 additions and 5 deletions

2
run.py
View File

@@ -3,6 +3,8 @@ import sys
from typing import Union, OrderedDict
sys.path.insert(0, os.getcwd())
# must come before ANY torch or fastai imports
import toolkit.cuda_malloc
import argparse
from toolkit.job import get_job