Commit 6cbab785 authored by novelailab's avatar novelailab

more fixing

parent 0879cf91
...@@ -58,7 +58,7 @@ class BasedOptimizer: ...@@ -58,7 +58,7 @@ class BasedOptimizer:
def step(self, scaler=None): def step(self, scaler=None):
if scaler: if scaler:
scaler.step(self.optimizer) scaler.step(self.optimizer)
else: else:
self.optimizer.step() self.optimizer.step()
...@@ -94,5 +94,6 @@ class BasedOptimizer: ...@@ -94,5 +94,6 @@ class BasedOptimizer:
def load(cls, path): def load(cls, path):
with open(path, 'rb') as f: with open(path, 'rb') as f:
based_optimizer = pickle.load(f) based_optimizer = pickle.load(f)
based_optimizer.optimizer.load_state_dict(torch.load(path)) based_optimizer.optimizer.load_state_dict(torch.load(path))
return based_optimizer return based_optimizer
\ No newline at end of file
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment