I notice that the model uses ReLU activations and therefore plain Xavier initialisation will be suboptimal for the main weights.
A better solution would be to supply the correct gain parameter for the activation.
https://pytorch.org/docs/stable/nn.html#torch.nn.init.calculate_gain 50
nn.init.xavier_uniform(m.weight.data, nn.init.calculate_gain('relu'))

With relu activation this almost gives you the Kaiming initialisation scheme. Kaiming uses either fan_in or fan_out, Xavier uses the average of fan_in and fan_out.
For the biases though, always initialising them to zero seems like a solid default choice.