Not working 50 series

[k95555]

Not working on RTX 5080, gives PyTorch Blackwell SM_120 error. Tried with the nightly version; training never started. I tried everything, wasted 4 days and 8 hours today, and I have given up.

[agarzon-orangerine]

I run it from docker ostris/aitoolkit:latest on my RTX 5090, no issues.

[zhimiants]

The same PyTorch Blackwell SM_120

[Norman3D]

pip uninstall torch torchvision torchaudio
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu128
pip install --no-build-isolation --pre -v -U git+https://github.com/facebookresearch/xformers.git@fde5a2fb46e3f83d73e2974a4d12caf526a4203e

[patrice74]

Generating baseline samples before training Error running job: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

======================================== Result:

• 0 completed jobs
• 1 failure ======================================== Traceback (most recent call last): File "D:\AI-Toolkit-Main\AI-Toolkit\run.py", line 120, in main() File "D:\AI-Toolkit-Main\AI-Toolkit\run.py", line 108, in main raise e File "D:\AI-Toolkit-Main\AI-Toolkit\run.py", line 96, in main job.run() File "D:\AI-Toolkit-Main\AI-Toolkit\jobs\ExtensionJob.py", line 22, in run process.run() File "D:\AI-Toolkit-Main\AI-Toolkit\jobs\process\BaseSDTrainProcess.py", line 1982, in run self.sample(self.step_num) File "D:\AI-Toolkit-Main\AI-Toolkit\extensions_built_in\sd_trainer\UITrainer.py", line 225, in sample super().sample(step, is_first) File "D:\AI-Toolkit-Main\AI-Toolkit\jobs\process\BaseSDTrainProcess.py", line 363, in sample self.sd.generate_images(gen_img_config_list, sampler=sample_config.sampler) File "D:\AI-Toolkit-Main\AI-Toolkit\venv\lib\site-packages\torch\utils_contextlib.py", line 120, in decorate_context return func(*args, **kwargs) File "D:\AI-Toolkit-Main\AI-Toolkit\toolkit\stable_diffusion_model.py", line 1147, in generate_images network.merge_in(merge_weight=merge_multiplier) File "D:\AI-Toolkit-Main\AI-Toolkit\toolkit\network_mixins.py", line 745, in merge_in module.merge_in(merge_weight) File "D:\AI-Toolkit-Main\AI-Toolkit\venv\lib\site-packages\torch\utils_contextlib.py", line 120, in decorate_context return func(*args, **kwargs) File "D:\AI-Toolkit-Main\AI-Toolkit\toolkit\network_mixins.py", line 374, in merge_in weight RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

Windows 11 + 5090 install AI-Toolkit-Easy-Install.bat and Pinokio the same error

[patrice74]

Norman3D's solution is the right one. Thank you. Windows 11 + 5090 Python 3.12.10 CUDA 12.8 PyTorch 2.8

git clone https://github.com/ostris/ai-toolkit.git

cd ai-toolkit
python -m venv venv
.\venv\Scripts\activate
pip install -r requirements.txt

cd ui
pip install python-dotenv
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu128
pip install --upgrade diffusers[torch]

Compile latest version of xformers (0.32)

git config --global core.longpaths true
pip install wheel setuptools packaging
pip install --no-build-isolation --pre -v -U git+https://github.com/facebookresearch/xformers.git@fde5a2fb46e3f83d73e2974a4d12caf526a4203e

or install a pre-compiled version of xformers (0.30) https://huggingface.co/Panchovix/xformers-windows-blackwell2.0-nightly/tree/main pip install path_to xformers-0.0.30%2B9a2cd3ef.d20250321-cp312-cp312-win_amd64.whl

triton installation pip install triton-windows

Thanks again Norman3D

[omgitsgb]

I have a one click installer I made just in case you guys were not able to get it working. Just have Python 3.12 installed.

Cheers!

https://github.com/omgitsgb/ostris-ai-toolkit-50gpu-installer

[Nightspeller]

This helped for me: pip install -U torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128