i am loading and using dropout in pretrained vovnet model(from timm package) like this :
class CassvaImgClassifier(nn.Module):
    def __init__(self, model_arch, n_class, pretrained=False):
        super().__init__()
        self.model = timm.create_model(model_arch, pretrained=pretrained)
        n_features = self.model.head.fc.in_features
        #self.model.head.fc = nn.Linear(n_features, n_class)
        
        self.model.head.fc = nn.Sequential(
            nn.Dropout(0.3),
            #nn.Linear(n_features, hidden_size,bias=True), nn.ELU(),
            nn.Linear(n_features, n_class, bias=True)
        )
        
    def forward(self, x):
        x = self.model(x)
        return x

i want to replace dropout with multi sample dropout : GitHub - lonePatient/multi-sample_dropout_pytorch: a simple pytorch implement of Multi-Sample Dropout 10
what should be the updated code?