Compute loss with subtensor of weight matrix?


I am trying to compute a loss using a selection of rows from a weight matrix. Is this possible? The reason I ask is 1) because I can’t quite grasp how indexing an output tensor and only using part of it for the loss affects the computational graph and backpropagation, and 2) I don’t know whether it’s possible to use parameters in the loss.

w = net.layer1.parameters()[0]
w_row = w[2,:]
w_target = torch.randn(w_rowshape)
loss = custom_criterion(w_row w_target)