Commit 2b7b9351 authored by Verena Praher's avatar Verena Praher
Browse files

add utility function to save models

parent e85bb786
......@@ -240,6 +240,13 @@ def preprocess_specs(source_root, destination_root, frame_length=256, hop=1.0):
np.save(destination, framed_spec)
def save(model, path):
try:
torch.save(model.module.state_dict(), path)
except AttributeError:
torch.save(model.state_dict(), path)
if __name__=='__main__':
# TESTS
......
Supports Markdown
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