def METHOD_NAME(x: torch.Tensor, expected_sharded_ten: ShardedTensor):
    pad_args = []
    assert len(x.shape) == len(expected_sharded_ten.local_shape)
    # Reversed iteration order because F.pad expects so
    for x_sh, exp_sh, axis_fragm in reversed(
        list(
            zip(x.shape, expected_sharded_ten.local_shape, expected_sharded_ten.axis_fragmentations)
        )
    ):
        if x_sh == exp_sh:
            pad_args.extend((0, 0))
        elif x_sh > exp_sh:
            assert (
                False
            ), f'Expected shape ({exp_sh}) smaller than actual ({x_sh}) for {repr(expected_sharded_ten)}'
        else:
            pad_args.extend((0, exp_sh - x_sh))
    # TODO: behavior control with envvar is for testing purposes only, remove it
    if not int(os.environ.get('DIST_CKPT_PAD_REPLICATE', 0)):
        return torch.nn.functional.pad(x, pad_args)
    # unsqueeze and squeeze to get shapes supported by cudnn
    print(f'Replicating last row for {expected_sharded_ten.key}')
    if x.dtype == torch.bfloat16:
        return (
            torch.nn.functional.pad(x.float().unsqueeze(0), pad_args, mode='replicate')
            .squeeze(0)
            .bfloat16()
        )
    return torch.nn.functional.pad(x.unsqueeze(0), pad_args, mode='replicate').squeeze(0)