Commit c2c88a53 authored by Andrei-Claudiu Roibu's avatar Andrei-Claudiu Roibu 🖥
Browse files

added save function

parent 4917a5c2
......@@ -103,10 +103,23 @@ class BrainMapperUNet(nn.Module):
return probability_map
def save(self, path):
"""Model Saver
Function saving the model with all its parameters to a given path.
The path must end with a *.model argument.
Args:
path (str): Path string
Returns:
None
Raises:
None
"""
Description
"""
pass
print("Saving Model... {}".format(path))
torch.save(self, path)
def predict(self, X):
"""
......
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