CUDA < 12.4 should work. 12.4+ works with some patches check this issue.
module load <cuda 11.8>
export CUDA_HOME=<path to cuda 11.8> # adjust this to your environment check by running 'module show <module_name>'
export PATH=$CUDA_HOME/bin:$PATH
export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH
| # Based on https://discuss.pytorch.org/t/check-gradient-flow-in-network/15063/10 | |
| def plot_grad_flow(named_parameters): | |
| '''Plots the gradients flowing through different layers in the net during training. | |
| Can be used for checking for possible gradient vanishing / exploding problems. | |
| Usage: Plug this function in Trainer class after loss.backwards() as | |
| "plot_grad_flow(self.model.named_parameters())" to visualize the gradient flow''' | |
| ave_grads = [] | |
| max_grads= [] |