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

added save function

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