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.