Added cuda transfer check
@@ -24,3 +24,5 @@ if __name__ == "__main__":
out = net(v)
print(out)
print("Cuda's availability is %s" % torch.cuda.is_available())
+ if torch.cuda.is_available():
+ print("Data from cuda: %s" % out.to('cuda'))
