diff --git a/py/test_gpu.py b/py/test_gpu.py
new file mode 100644
index 0000000..a06d44f
--- /dev/null
+++ b/py/test_gpu.py
@@ -0,0 +1,12 @@
+#%%
+%lsmagic
+
+#%%
+import torch
+t_cpu = torch.rand(500,500,500)
+%timeit t_cpu @ t_cpu
+
+t_gpu = torch.rand(500,500,500).cuda()
+%timeit t_gpu @ t_gpu
+
+#TODO add TensorFlow GPU check
\ No newline at end of file