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

added test for parameters being on stored on GPU

parent c2c88a53
...@@ -120,9 +120,27 @@ class BrainMapperUNet(nn.Module): ...@@ -120,9 +120,27 @@ class BrainMapperUNet(nn.Module):
print("Saving Model... {}".format(path)) print("Saving Model... {}".format(path)), path), path)
def test_if_cuda(self):
"""Cuda Test
This function tests if the model parameters are allocated to a CUDA enabled GPU.
bool: Flag indicating True if the tensor is stored on the GPU and Flase otherwhise
return next(self.parameters()).is_cuda
def predict(self, X): def predict(self, X):
""" """Post-training Output Prediction
Description Description
""" """
return None return None
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