Speed up differentiable 5PC and fix the batch size issue #3062

Job	Run time
${{ matrix.os }} - torch-nightly, ${{ matrix.pytorch-dtype }}	0s
coverage / dynamo (float32)	1m 10s
coverage / slow (float32)	1m 13s
coverage / overall (float32)	1m 12s
coverage / slow (float64)	1m 10s
coverage / overall (float64)	1m 11s
docs / docstests & sphinx-build - python-3.11, torch-2.3.0	1m 9s
tests-cpu-macos / python-3.8, torch-2.2.1	1m 12s
tutorials / python-3.11, torch-2.3.0	1m 10s
typing / python-3.11, torch-2.3.0	58s
tests-cpu-macos / python-3.11, torch-2.2.1	1m 14s
tests-cpu (Ubuntu-latest, float32) / python-3.8, torch-1.9.1	1m 8s
tests-cpu (Ubuntu-latest, float32) / python-3.8, torch-2.3.0	1m 11s
tests-cpu (Ubuntu-latest, float64) / python-3.8, torch-1.9.1	1m 11s
tests-cpu (Ubuntu-latest, float32) / python-3.11, torch-2.3.0	1m 12s
tests-cpu (Ubuntu-latest, float64) / python-3.8, torch-2.3.0	1m 8s
tests-cpu (Ubuntu-latest, float64) / python-3.11, torch-2.3.0	1m 10s
tests-cpu (Windows-latest, float32) / python-3.8, torch-1.9.1	56s
tests-cpu (Windows-latest, float64) / python-3.8, torch-1.9.1	1s
tests-cpu (Windows-latest, float32) / python-3.8, torch-2.3.0	1m 4s
tests-cpu (Windows-latest, float64) / python-3.8, torch-2.3.0	1s
tests-cpu (Windows-latest, float32) / python-3.11, torch-2.3.0	1m 6s
tests-cpu (Windows-latest, float64) / python-3.11, torch-2.3.0	1m 3s
coverage / report	1s
collector	1s
	22m 52s

Provide feedback