@smonsays overall it looks a bit overcomplicated, my comment re the changing probabilities issue, I’d do it one of two ways

encapsulate the DropConnect in its own module that has the drop prob as a scalar attribute, iterate over the model modules when you want to change it, and adjust that scalar for your DropConnect modules.
rely on Python pass-by-reference and wrap the scalar drop value in a class with a float conversion and call float(drop_prob)  when you want to use it, you can impl your drop connect as functional or module

In both of the above, I’d use torch.empty(tensor_size).bernouilli_(scalar_prob) instead of the torch.bernoulli which requires a tensor for prob arg and thus you must match device, etc . Also, I don’t think prob should be a buffer, it will end up in the state dict, it is a training detail in most situations.