Fix TransformerEngine and Activation checkpointing interaction for thunderFX #5033

Job	Run time
pytester (ubuntu-22.04, 3.10, latest, core)	12m 15s
pytester (ubuntu-22.04, 3.10, latest, ops)	27m 14s
pytester (ubuntu-22.04, 3.10, nightly, core)	11m 56s
pytester (ubuntu-22.04, 3.10, nightly, ops)	26m 48s
pytester (macOS-14, 3.10, latest, core)	10m 0s
pytester (macOS-14, 3.10, latest, ops)	23m 29s
pytester (macOS-14, 3.10, nightly, core)	10m 10s
pytester (macOS-14, 3.10, nightly, ops)	20m 12s
pytester (windows-2022, 3.10, latest, core)	23m 58s
pytester (windows-2022, 3.10, nightly, core)	26m 7s
pytester (ubuntu-22.04, 3.11, latest)	23m 28s
pytester (ubuntu-22.04, 3.12, latest)	27m 59s
testing-guardian	1s
	4h 3m 37s

Provide feedback