Skip to content

Commit 1dbce17

Browse files
authored
Improve pytorch unit test (#2290)
1 parent b2226eb commit 1dbce17

File tree

1 file changed

+7
-0
lines changed

1 file changed

+7
-0
lines changed

tests/by_image/pytorch-notebook/units/unit_pytorch.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,3 +3,10 @@
33
import torch
44

55
print(torch.tensor([[1.0, 4.0, 7.0], [4.0, 9.0, 11.0]]))
6+
7+
# Check if GPU is available
8+
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
9+
print(f"Using device: {device}")
10+
# Create large tensors
11+
size = 10000
12+
torch.randn(size, size, device=device)

0 commit comments

Comments
 (0)