Skip to content

Add a workflow to run benchmarks #6

Add a workflow to run benchmarks

Add a workflow to run benchmarks #6

Workflow file for this run

name: Benchmarks
on:
push:
branches:
- main
pull_request:
branches:
- main
jobs:
multiprocess_gpu_test:
runs-on: ubuntu-latest # For self-hosted runners with GPUs, update accordingly.
strategy:
matrix:
gpu-count: [2, 4, 8] # Adjust based on the runner's available GPUs
steps:
- uses: actions/checkout@v3
- name: Install dependencies
run: |
python -m pip install --upgrade pip
pip install -U jax
pip install -U "jax[cuda12]"
pip install pytest
- name: Run Multiprocess GPU Test
run: |
python -m pytest tests/multiprocess_gpu_test.py