How to random the weight of the first layer at each training step

[chao0716]

Dear,

I am from the robot area. I am trying to random the weight of the first layer at each training step. it`s a method from a domain randomization paper.

For example this simple code in pytorch lighting doc.
https://colab.research.google.com/drive/1rHBxrtopwtF8iLpmC_e7yl3TeDGrseJL?usp=sharing%3E#scrollTo=V7ELesz1kVQo
at each training step, I hope random the set the weights of the first layer : nn.Linear(28 * 28, 128), nn.ReLU(), nn.Linear(128, 3))

How could I do it in pytorch lightining framework?

Thanks for your help