pynvml is not supported on GeForce

I am encountering an issue when attempting to use the triangle_multiplicative_update function from the cuequivariance-torch library. Despite triton==3.3.0 was installed as recommended by the traceback, the function still fails to import.

Environment

GPU and CUDA version

• NVIDIA GeForce RTX 3090
• CUDA 12.6 (installed from conda)

conda environment

channels:
  - conda-forge
  - nvidia/label/cuda-12.6.0
dependencies:
  - nvidia/label/cuda-12.6.0::cuda-toolkit
  - python=3.12
  - pip:
    - cuequivariance-ops-torch-cu12==0.6.0
    - cuequivariance-torch==0.6.0
    - torch==2.7.0+cu126
    - torchaudio==2.7.0+cu126
    - torchvision==0.22.0+cu126

Minimal Reproducible code

import torch
import triton
from cuequivariance_torch import triangle_multiplicative_update

print(f"PyTroch vresion: {torch.__version__}\nTriton version: {triton.__version__}")
if torch.cuda.is_available():
    device = torch.device("cuda")
    batch_size, seq_len, hidden_dim = 1, 128, 128
    # Create input tensor
    x = torch.randn(batch_size, seq_len, seq_len, hidden_dim, requires_grad=True, device=device)
    # Create mask (1 for valid positions, 0 for masked)
    mask = torch.ones(batch_size, seq_len, seq_len, device=device)
    # Perform triangular multiplication
    output = triangle_multiplicative_update(
        x=x,
        direction="outgoing",  # or "incoming"
        mask=mask,
    )
    print(output.shape)  # torch.Size([1, 128, 128, 128])
    # Create gradient tensor and perform backward pass
    grad_out = torch.randn_like(output)
    output.backward(grad_out)
    # Access gradients
    print(x.grad.shape)

Output

PyTroch vresion: 2.7.0+cu126
Triton version: 3.3.0
Traceback (most recent call last):
  File "/path/to/cuequivariance-test/test.py", line 14, in <module>
    output = triangle_multiplicative_update(
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/path/to/env/lib/python3.12/site-packages/cuequivariance_torch/primitives/triangle.py", line 231, in triangle_multiplicative_update
    return f(
           ^^
  File "/path/to/env/lib/python3.12/site-packages/cuequivariance_ops_torch/__init__.py", line 72, in triangle_multiplicative_update
    raise Exception(
Exception: Failed to import Triton-based component: triangle_multiplicative_update:
Not Supported
Please make sure to install triton==3.3.0. Other versions may not work!

Start with cuequivariance_torch/primitives/triangle.py at the triangle_multiplicative_update call and cuequivariance_ops_torch/init.py line 72, then run the provided RTX 3090 reproduction with the stated package versions. Done means the reported compatibility behavior is resolved or the supported hardware and configuration are clearly established.