3D Flex pytorch issue on RTX4090

3D Flexible Refinement
13

@alexwoo1221 @Rafa @UCBKurt We have confirmed problems running 3DFlex jobs on a Hopper GPU and are investigating a possible fix.