It the variational optimisation something like a weight initialization scheme?
If so, you could create a function with your code and use apply on your model.
def weights_init(m):
    if isinstance(m, nn.Conv2d):
        "Your code here e.g.:"
        xavier(m.weight.data)
        xavier(m.bias.data)

model.apply(weights_init)