Skip to content

ai

Flax

🔗 https://flax.readthedocs.io

Package description

Flax is a high-performance neural network library and ecosystem for JAX that is designed for flexibility: Try new forms of training by forking an example and by modifying the training loop, not by adding features to a framework.

Use latest version

module load Flax

Use specific version

module load Flax/0.8.4-gfbf-2023a

PyTorch

🔗 https://pytorch.org/

Package description

Tensors and Dynamic neural networks in Python with strong GPU acceleration. PyTorch is a deep learning framework that puts Python first.

Use latest version

module load PyTorch

Use specific version

module load PyTorch/1.12.0-foss-2022a-CUDA-11.7.0
module load PyTorch/1.12.0-foss-2022a

PyTorch-bundle

🔗 https://pytorch.org/

Package description

PyTorch with compatible versions of official Torch extensions.

Use latest version

module load PyTorch-bundle

Use specific version

module load PyTorch-bundle/1.12.1-foss-2022a-CUDA-11.7.0

Uni-Core

🔗 https://github.com/dptech-corp/Uni-Core

Package description

An efficient distributed PyTorch framework

Use latest version

module load Uni-Core

Use specific version

module load Uni-Core/0.0.3-foss-2022a-CUDA-11.7.0

colossalai

🔗 https://colossalai.org/

Package description

Colossal-AI: A Unified Deep Learning System for Big Model Era

Use latest version

module load colossalai

Use specific version

module load colossalai/0.1.8-foss-2021a-CUDA-11.3.1

jax

🔗 https://jax.readthedocs.io/

Package description

Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more

Use latest version

module load jax

Use specific version

module load jax/0.4.25-gfbf-2023a