diff --git a/tests/by_image/pytorch-notebook/units/unit_pytorch.py b/tests/by_image/pytorch-notebook/units/unit_pytorch.py
index 1b739a59..eebc4d3d 100644
--- a/tests/by_image/pytorch-notebook/units/unit_pytorch.py
+++ b/tests/by_image/pytorch-notebook/units/unit_pytorch.py
@@ -3,3 +3,10 @@
 import torch
 
 print(torch.tensor([[1.0, 4.0, 7.0], [4.0, 9.0, 11.0]]))
+
+# Check if GPU is available
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Using device: {device}")
+# Create large tensors
+size = 10000
+torch.randn(size, size, device=device)