mrq/DL-Art-School
1
3
Fork 5
Code Issues Pull Requests Packages Projects Releases Wiki Activity

I guess I can't easily toggle it outside of here, but it works

Browse Source
This commit is contained in:
mrq 2023-02-23 07:02:06 +00:00
parent 0ef8ab6872
commit 7bcedca771
1 changed files with 14 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
codes/torch_intermediary/__init__.py
View File

@ -11,11 +11,25 @@ from torch.optim.adam import Adam
from torch.optim.adamw import AdamW from torch.optim.adamw import AdamW
""" """
"""
OVERRIDE_LINEAR = False OVERRIDE_LINEAR = False
OVERRIDE_EMBEDDING = False OVERRIDE_EMBEDDING = False
OVERRIDE_ADAM = False # True OVERRIDE_ADAM = False # True
OVERRIDE_ADAMW = False # True OVERRIDE_ADAMW = False # True
"""
USE_STABLE_EMBEDDING = True USE_STABLE_EMBEDDING = True
try:
import bitsandbytes as bnb
OVERRIDE_LINEAR = False
OVERRIDE_EMBEDDING = False
OVERRIDE_ADAM = True
OVERRIDE_ADAMW = True
except Exception as e:
OVERRIDE_LINEAR = False
OVERRIDE_EMBEDDING = False
OVERRIDE_ADAM = False
OVERRIDE_ADAMW = False
if OVERRIDE_LINEAR: if OVERRIDE_LINEAR:
from bitsandbytes.nn import Linear8bitLt as Linear from bitsandbytes.nn import Linear8bitLt as Linear