`pixi.toml` file for pytorch

pixi.toml

[project]
name = "my-ml-project"
description = "A project that does ML stuff"
authors = ["Your Name <your.name@gmail.com>"]
channels = ["conda-forge"]
platforms = ["linux-64", "osx-arm64"]

[tasks]
train-model = "python train.py"
evaluate-model = "python test.py"

[dependencies]
python = "3.11.*"
pytorch = ">=2.0.1"
torchvision = {version = ">=0.15"}
polars = ">=0.20,<0.21"
matplotlib-base = ">=3.8.2,<3.9"
ipykernel = ">=6.28.0,<6.29"

[feature.cuda]
platforms = ["linux-64"]
channels = ["conda-forge"]
system-requirements = {cuda = "12"}

[feature.cuda.tasks]
train-model = "python train.py --cuda"
evaluate-model = "python test.py --cuda"

[feature.cuda.dependencies]
pytorch-gpu = ">=2.0.1"
transformers = ">=4.39.3,<4.40"
jupyterlab = ">=4.1.5,<4.2"

[feature.mlx]
platforms = ["osx-arm64"]

[feature.mlx.tasks]
train-model = "python train.py --mlx"
evaluate-model = "python test.py --mlx"

[feature.mlx.dependencies]
pytorch = ">=2.1.2,<2.2"
transformers = ">=4.39.3,<4.40"
jupyterlab = ">=4.1.5,<4.2"

[feature.cpu]
platforms = ["linux-64", "osx-arm64"]

[environments]
cuda = ["cuda"]
mlx = ["mlx"]
default = ["cpu"]