Fix pAE-related bugs #2185

	name: Pytest
	on: [push, pull_request]

	env:
	TYPECHECK: True
	DEBUG: True
	DEEPSPEED_CHECKPOINTING: False
	IS_GITHUB_CI: True

	jobs:
	build:

	runs-on: ubuntu-latest
	timeout-minutes: 20
	strategy:
	fail-fast: false
	matrix:
	group: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

	steps:
	- uses: actions/checkout@v4
	- name: Set up Python 3.10
	uses: actions/setup-python@v5
	with:
	python-version: "3.10"
	- name: Install Nim
	uses: jiro4989/setup-nim-action@v2
	with:
	nim-version: '2.0.8'
	repo-token: ${{ secrets.GITHUB_TOKEN }}
	- name: Install Nim dependencies
	run: \|
	nimble install
	- name: Install dependencies
	run: \|
	python -m pip install uv
	python -m uv pip install --upgrade pip
	python -m uv pip install torch==2.3.1
	python -m uv pip install -e .[test]
	- name: Test with pytest
	run: \|
	python -m pytest --num-shards 10 --shard-id ${{ matrix.group }} tests/

Provide feedback