add torch rocm support to env

This commit is contained in:
2022-11-13 14:26:04 +01:00
parent 3727c74c8f
commit 7fc87def1e

View File

@ -14,15 +14,45 @@
overlays = [ devshell.overlay ];
};
# TODO: Originally it was nixpkgs.fetchurl but that didn't work, pkgs.fetchurl did...
# Determine the difference between nixpkgs and pkgs
# Taken from: https://github.com/gbtb/nix-stable-diffusion/blob/master/flake.nix
# Overlay: https://nixos.wiki/wiki/Overlays
# FetchURL: https://ryantm.github.io/nixpkgs/builders/fetchers/
torch-rocm = pkgs.hiPrio (pkgs.python310Packages.torch-bin.overrideAttrs (old: {
src = pkgs.fetchurl {
name = "torch-1.12.1+rocm5.1.1-cp310-cp310-linux_x86_64.whl";
url = "https://download.pytorch.org/whl/rocm5.1.1/torch-1.12.1%2Brocm5.1.1-cp310-cp310-linux_x86_64.whl";
hash = "sha256-kNShDx88BZjRQhWgnsaJAT8hXnStVMU1ugPNMEJcgnA=";
};
}));
torchvision-rocm = pkgs.hiPrio (pkgs.python310Packages.torchvision-bin.overrideAttrs (old: {
src = pkgs.fetchurl {
name = "torchvision-0.13.1+rocm5.1.1-cp310-cp310-linux_x86_64.whl";
url = "https://download.pytorch.org/whl/rocm5.1.1/torchvision-0.13.1%2Brocm5.1.1-cp310-cp310-linux_x86_64.whl";
hash = "sha256-mYk4+XNXU6rjpgWfKUDq+5fH/HNPQ5wkEtAgJUDN/Jg=";
};
}));
myPython = pkgs.python310.withPackages (p: with p; [
# Basic
rich
# Discord
discordpy
python-dotenv
pynacl
# Scraping
beautifulsoup4
requests
pynacl
rich
torch
# MachineLearning
torch-rocm
torchvision-rocm
numpy
matplotlib
nltk
]);
in {