AttributeError: 'torchvision' has no attribute 'extension' When Using Unsloth on Kaggle

[Saber120]

I am trying to train a model using the Unsloth library on Kaggle, but I am encountering the following error during the import process:

AttributeError: partially initialized module 'torchvision' has no attribute 'extension' (most likely due to a circular import)

Here are the steps I have taken in my environment setup:

%%capture
!mamba install --force-reinstall aiohttp -y
!pip install -U "xformers<0.0.26" --index-url https://download.pytorch.org/whl/cu121
!pip install "unsloth[kaggle-new] @ git+https://github.com/unslothai/unsloth.git"

# Temporary fix for https://github.com/huggingface/datasets/issues/6753
!pip install datasets==2.16.0 fsspec==2023.10.0 gcsfs==2023.10.0

import os
os.environ["WANDB_DISABLED"] = "true"

from unsloth import FastLanguageModel
import torch
max_seq_length = 2048
dtype = None
load_in_4bit = True

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "unsloth/SmolLM2-360M-Instruct-bnb-4bit",
    max_seq_length = max_seq_length,
    dtype = dtype,
    load_in_4bit = load_in_4bit,
)

Can anyone help me troubleshoot this issue? Thank you!

[Erland366]

Can you try using Unsloth's new way of installation in Kaggle?

%%capture
!pip install pip3-autoremove
!pip-autoremove torch torchvision torchaudio -y
!pip install torch torchvision torchaudio xformers --index-url https://download.pytorch.org/whl/cu121
!pip install unsloth
!pip install --upgrade --no-cache-dir git+https://github.com/huggingface/transformers.git git+https://github.com/huggingface/trl.git