diff options
author | Tim Dettmers <tim.dettmers@gmail.com> | 2021-11-29 08:24:17 -0800 |
---|---|---|
committer | Tim Dettmers <tim.dettmers@gmail.com> | 2021-11-29 08:24:17 -0800 |
commit | 3cff6795fb70dd99b4802593f3c70d291e0cd1dc (patch) | |
tree | 0a20b7bb88e1db4be688ee148fd957eeee0dc537 /bitsandbytes/optim | |
parent | 108cf9fc1f8c6bc0360a49ce790699928883b3d3 (diff) | |
parent | 262350c10f4174d0c775b61be2dbf526afa69cd2 (diff) |
Merge branch 'main' of github.com:facebookresearch/bitsandbytes into 0.26.0
Diffstat (limited to 'bitsandbytes/optim')
-rw-r--r-- | bitsandbytes/optim/adam.py | 11 |
1 files changed, 8 insertions, 3 deletions
diff --git a/bitsandbytes/optim/adam.py b/bitsandbytes/optim/adam.py index 1e93a60..ed1b9f0 100644 --- a/bitsandbytes/optim/adam.py +++ b/bitsandbytes/optim/adam.py @@ -2,7 +2,12 @@ # # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. + +import math +import os + import torch +import torch.distributed as dist from bitsandbytes.optim.optimizer import Optimizer2State import bitsandbytes.functional as F @@ -219,9 +224,9 @@ class AnalysisAdam(torch.optim.Optimizer): if self.savedir != '' and state['step'] % 100 == 0: if not os.path.exists(self.savedir): os.makedirs(self.savedir) shapestr = '_'.join([str(dim) for dim in p_data_fp32.shape]) - pathe = join(self.savedir, f'{p_id}_{shapestr}_abserr.pkl') - pathrele = join(self.savedir, f'{p_id}_{shapestr}_relerr.pkl') - pathcounts = join(self.savedir, f'{p_id}_{shapestr}_counts.pkl') + pathe = os.path.join(self.savedir, f'{p_id}_{shapestr}_abserr.pkl') + pathrele = os.path.join(self.savedir, f'{p_id}_{shapestr}_relerr.pkl') + pathcounts = os.path.join(self.savedir, f'{p_id}_{shapestr}_counts.pkl') torch.save(e, pathe) torch.save(rele, pathrele) torch.save(counts, pathcounts) |