def METHOD_NAME(self, destination=None, prefix="", keep_vars=False):
    if self.use_sharded_state:
        return super().local_state_dict(
            destination=destination, prefix=prefix, keep_vars=keep_vars
        )
    else:
        if self.rank == 0:
            return super().METHOD_NAME(
                destination=destination, prefix=prefix, keep_vars=keep_vars
            )
        else:
            # We must call state_dict() due to use of communication
            # primitives. But we don't use the result.
            super().METHOD_NAME()
            return destination or {}