bug: tests failing at nvidia-driver-530

[christallire]

Description

I was setting up a fresh pytorch environment and the latest cuda package installs nvidia-driver-530 and it makes kernl tests fail
and prints RuntimeError: CUDA error: misaligned address

also optimize_model at inference time prints the same error message too so I had to comment it out.

Steps to reproduce

Prepare fresh ubuntu environment

1. sudo apt install nvidia-drivers-520
2. sudo apt install cuda=11.7.1-1
3. pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu117
4. cd ~/kernl/ && pytest

Expected Behavior

tests pass

Actual Behavior

=================================================================================== test session starts ===================================================================================
platform linux -- Python 3.10.9, pytest-7.2.1, pluggy-1.0.0
rootdir: /home/chris/workspace/kernl
collected 2883 items                                                                                                                                                                      

test/test_attention.py ..FEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEE [  5%]
EEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEEE

Your environment

• Ubuntu 22 LTS
• Linux d42gx-1 5.19.0-35-generic Feat/layernorm #36-Ubuntu SMP PREEMPT_DYNAMIC Fri Feb 3 18:36:56 UTC 2023 x86_64 x86_64 x86_64 GNU/Linux
• Python 3.10.9
• pip 23.0.1 from /home/chris/mambaforge/envs/singu/lib/python3.10/site-packages/pip (python 3.10)

Self-service

• I would be willing to help fix this bug myself.

Code of Conduct

• I agree to follow this project's Code of Conduct