41 lines
1.3 KiB
Nix
41 lines
1.3 KiB
Nix
with import <nixpkgs> {
|
|
config.sudaSupport = true;
|
|
config.allowUnfree = true;
|
|
};
|
|
let
|
|
python3 = pkgs.python311;
|
|
in
|
|
mkShell {
|
|
name = "tensorflow-cuda-shell";
|
|
buildInputs = with python3.pkgs; [
|
|
pip
|
|
numpy
|
|
setuptools
|
|
virtualenv
|
|
tensorflow tensorboard
|
|
];
|
|
shellHook = ''
|
|
export CUDA_PATH=${pkgs.cudatoolkit}
|
|
export LD_LIBRARY_PATH=${pkgs.stdenv.cc.cc.lib}/lib:${pkgs.cudaPackages_11.cudatoolkit}/lib:${pkgs.cudaPackages_11.cudnn}/lib:${pkgs.cudaPackages_11.cudatoolkit.lib}/lib:$LD_LIBRARY_PATH
|
|
export EXTRA_LDFLAGS="-L/lib -L${pkgs.linuxPackages.nvidia_x11}/lib"
|
|
export EXTRA_CCFLAGS="-I/usr/include"
|
|
alias pip="PIP_PREFIX='$(pwd)/_build/pip_packages' TMPDIR='$HOME' \pip"
|
|
export PYTHONPATH="$(pwd)/_build/pip_packages/lib/python3.7/site-packages:$PYTHONPATH"
|
|
export PATH="$(pwd)/_build/pip_packages/bin:$PATH"
|
|
unset SOURCE_DATE_EPOCH
|
|
|
|
# Use bash instead of zsh
|
|
export SHELL=${pkgs.bash}/bin/bash
|
|
|
|
# set up a virtuial environment
|
|
export VIRTUAL_ENV=$PWD/venv
|
|
if [ ! -d "$VIRTUAL_ENV" ]; then
|
|
virtualenv "$VIRTUAL_ENV"
|
|
source "$VIRTUAL_ENV/bin/activate"
|
|
python -m pip install --upgrade pip
|
|
# python -m pip install -r requirements.txt
|
|
else
|
|
source "$VIRTUAL_ENV/bin/activate"
|
|
fi
|
|
'';
|
|
} |