Hi,
I’d write it in a (maybe verbose) way like
class Net(nn.Module):
    def __init__(self, R):
        super().__init__()
        self.lin1 = nn.Linear(2, 100)
        self.lin2 = nn.Linear(100, 3)

        with torch.no_grad():
            weight = self.lin2.weight.data.clone().detach()
            weight.mul_( (R / torch.sqrt(torch.sum(weight**2 ,axis=1)))[:, None] )
            self.lin2.weight.data.copy_(weight)
            del weight

    def forward(self, x):
        return self.lin2(F.relu(self.lin1(x)))

A toy colab is here: https://colab.research.google.com/drive/12yqnYoRF1C_of6s8fokTEqK2h9wl-wvj?usp=sharing 1