@ptrblck For one of my models I am using a pytorch loss function. I do some calculations before hand but in the end the loss is being handled by pytorch.
a = torch.sum(outputs * one_hot)
F.mse_loss(a, desired).backward(retain_graph=True)
self.optimizer.step()

For another model I implemented the loss function:
loss = -(obj_func + tau * entropy)
loss.backward(retain_graph=True)
self.optimizer.step()

Yet the results for both are the same, the last layer and layer normalizations are updating but the remaining are not with the gradient distribution mentioned above.
I am trying to use kaiming and xavier initialization, but the initialized weights are not saving to the model. The code I am using to try accomplish this:
new_weights = collections.OrderedDict()
for key, value in network.state_dict().items():
    try:
            if 'bias' in key:
                new_weights[key] = torch.full_like(network.state_dict()[key], 0)
            elif 'ltsm' in key:
                new_weights[key] = torch.nn.init.xavier_uniform_(value, gain=5/3)
            elif 'batch' in key or 'norm' in key:
                new_weights[key] = torch.nn.init.xavier_uniform_(value)
                # Is the right initialization to use for batch and layer normalization?
            else:
                new_weights[key] = torch.nn.init.kaiming_uniform_(value, a=0, mode='fan_out', nonlinearity = 'relu')
    except:
        new_weights[key] = torch.full_like(network.state_dict[key], 0)

I initially tried to directly reassign the weights but that did not work.