def METHOD_NAME(step_idx, num_steps):
    if step_idx < num_steps - 1 or not unfreeze_params:
        # If the input does not require gradient, then the 0th
        # frozen linear gets resharded in the catch-all reshard
        # since we cannot register an autograd hook on it
        expected_post_backward_reshard_count = (
            self.NUM_LINEARS if inp_requires_grad else self.NUM_LINEARS - 1
        )
    else:
        # This follows the normal post-backward hook path
        expected_post_backward_reshard_count = self.NUM_LINEARS
    self.assertEqual(
        post_backward_reshard_count, expected_post_backward_reshard_count
    )