(ak-uproot) ➜  ~ echo $LD_LIBRARY_PATH
/usr/local/cuda-12.6/lib64:/home/iason/software/torch/nccl/build/lib
(ak-uproot) ➜  ~ ipython
iPython 3.12.9 | packaged by conda-forge | (main, Feb 14 2025, 08:00:06) [GCC 13.3.0]
Type 'copyright', 'credits' or 'license' for more information
IPython 8.32.0 -- An enhanced Interactive Python. Type '?' for help.

In [1]: import torch

In [2]: torch.cuda.is_available()
Out[2]: True

In [3]: torch.cuda.is_bf16_supported()
Out[3]: True

In [4]: x = torch.rand(5, 3)

In [5]: print(x)
tensor([[0.3051, 0.5993, 0.2161],
        [0.1847, 0.2640, 0.4210],
        [0.0208, 0.4589, 0.3539],
        [0.5538, 0.0826, 0.1458],
        [0.3174, 0.1281, 0.6403]])

In [6]: import torch
   ...: import torch.nn as nn
   ...:
   ...: # Check if CUDA is available
   ...: device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
   ...: print(f"Using device: {device}")
   ...:
   ...: # Set parameters
   ...: batch_size = 8
   ...: seq_length = 10
   ...: embed_dim = 512
   ...: num_heads = 8
   ...:
   ...: # Create random input tensors
   ...: # Shape: [sequence length, batch size, embedding dimension]
   ...: query = torch.randn(seq_length, batch_size, embed_dim, device=device)
   ...: key = torch.randn(seq_length, batch_size, embed_dim, device=device)
   ...: value = torch.randn(seq_length, batch_size, embed_dim, device=device)
   ...:
   ...: # Create MultiheadAttention module
   ...: mha = nn.MultiheadAttention(embed_dim, num_heads).to(device)
   ...:
   ...: # Forward pass
   ...: output, attn_weights = mha(query, key, value)
   ...:
   ...: # Print shapes
   ...: print(f"Input shapes: query={query.shape}, key={key.shape}, value={value.shape}")
   ...: print(f"Output shape: {output.shape}")
   ...: print(f"Attention weights shape: {attn_weights.shape}")
Using device: cuda
Input shapes: query=torch.Size([10, 8, 512]), key=torch.Size([10, 8, 512]), value=torch.Size([10, 8, 512])
Output shape: torch.Size([10, 8, 512])
Attention weights shape: torch.Size([8, 10, 10])

In [7]: query.device
Out[7]: device(type='cuda', index=0)

In [8]: exit()
(ak-uproot) ➜  ~ pip list | grep torch
torch                     2.6.0+cu126
torchaudio                2.6.0+cu126
torchvision               0.21.0+cu126
(ak-uproot) ➜  ~ pip list | grep nvidia
(ak-uproot) ➜  ~