hazza/docker-stacks
1
0
Fork 0
mirror of https://github.com/jupyter/docker-stacks.git synced 2025-10-15 14:02:58 +00:00
Code Issues Packages Projects Releases Wiki Activity

Improve pytorch unit test (#2290)

Browse Source
This commit is contained in:
Ayaz Salikhov
2025-04-25 01:39:47 +01:00
committed by GitHub
parent b2226ebf0e
commit 1dbce17442
1 changed files with 7 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
tests/by_image/pytorch-notebook/units/unit_pytorch.py
View File

@@ -3,3 +3,10 @@
import torch import torch
print(torch.tensor([[1.0, 4.0, 7.0], [4.0, 9.0, 11.0]])) 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)