fix #2717

	name: CI

	on: [push, pull_request]

	jobs:
	pytest:
	runs-on: ubuntu-latest
	steps:
	- uses: actions/checkout@v2
	- name: Set up Python 3.9
	uses: actions/setup-python@v2
	with:
	python-version: 3.9
	- name: Install pytest
	run: \|
	python -m pip install --upgrade pip
	pip install pytest
	pip install .[jax_cpu]
	pip install .[pytorch_cpu]
	pip install .[wmt]
	pip install .[ogbg]
	- name: Run pytest
	run: \|
	pytest -vx tests/version_test.py
	pytest -vx tests/workloads/imagenet_resnet/imagenet_jax/workload_test.py
	pytest -vx tests/test_num_params.py
	pytest -vx tests/test_param_shapes.py

Provide feedback