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):
print("Saving Model... {}".format(path))
torch.save(self, path)
@property
def test_if_cuda(self):
"""Cuda Test
This function tests if the model parameters are allocated to a CUDA enabled GPU.
Args:
None
Returns:
bool: Flag indicating True if the tensor is stored on the GPU and Flase otherwhise
Raises:
None
"""
return next(self.parameters()).is_cuda
def predict(self, X):
"""
"""Post-training Output Prediction
Description
"""
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