training_with_trainable_dropout.py

#example training w/ trainable dropout layer
for epoch in range(epochs):
    for (train_batch, test_batch) in zip(train_dataloader, test_dataloader):
        features, target = train_batch
        test_features, test_target= test_batch
    
        unfreeze_all_but_dropout(model);
        
        optimizer.zero_grad();

        pred = model(features);
        loss = criterion_loss(pred,target);
        loss.backward();
        optimizer.step();
    
        freeze_all_but_dropout(model);

        test_pred = model(test_features);
        loss_test = criterion_loss(test_pred,test_target);
        loss_test.backward();
        optimizer.step();
        
        
# then your model would have a trainable dropout layer as its inital layer that looks like so:
class TrainableDropoutLayer(nn.Module):                                                                                                
    def __init__(self, num_features):                                                                                                  
        super(TrainableDropoutLayer, self).__init__()                                                                                  
                                                                                                                                       
        self.num_features = num_features                                                                                               
        self.scale = nn.Parameter(torch.ones(num_features))                                                                            
                                                                                                                                       
    def forward(self, x):                                                                                                              
        scale = torch.sigmoid(self.scale)                                                                                              
                                                                                                                                       
        return x * scale