Weight becomes untrainable after multiplying with constant

System information

  • Have I written custom code (as opposed to using a stock example script
    provided in TensorFlow)
    : yes
  • OS Platform and Distribution (e.g., Linux Ubuntu 16.04): macos
  • TensorFlow installed from (source or binary):
  • TensorFlow version (use command below): tensorflow 2.4
  • Python version: python 3.7

Describe the problem

I want to mask the weights of my network (set them to zero and freeze them, i.e. no gradients). I tried creating a custom layer and multiply the weight matrices by a constant matrix, but this makes the weight matrices no longer trainable. Note that this approach used to work in tensorflow 1, and I can’t figure out why it doesn’t work anymore in tf2. I’ve attached an example in this colab.

1 possible answer(s) on “Weight becomes untrainable after multiplying with constant