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

added a device flag to predict

parent 89a2e5fb
...@@ -139,7 +139,7 @@ class BrainMapperUNet(nn.Module): ...@@ -139,7 +139,7 @@ class BrainMapperUNet(nn.Module):
return next(self.parameters()).is_cuda return next(self.parameters()).is_cuda
def predict(self, X): def predict(self, X, device= 0):
"""Post-training Output Prediction """Post-training Output Prediction
This function predicts the output of the of the U-net post-training This function predicts the output of the of the U-net post-training
......
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