Hi
Why would your weights include a 3rd (batch_size?) dimension? Your state, which you show does and should, but for the weights that seems unusual.
That said, you can mostly assign to parameter’s .data submatrices.
Best regards
Thomas