[Bug Fix] Fix variable reference as key error in distributed shampoo

[c0derzer0]

This PR addresses the issue #285

While applying gradients with distributed shampoo, it fails with the error "TypeError: Variable is unhashable. Instead, use variable.ref() as the key".

This is caused because of using variable instead of variable.ref() as key which is fixed in this pull request. This has been a recurring issue whenever the lingvo library is reinstalled and it needs to be changed locally. This PR will also address this issue for any other users trying to use distributed shampoo.

[google-cla[bot]]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.