def test_modify_dropout(net, img_size):
    """Test ability to modify dropout and fc modules of network"""
    dropout = nn.Sequential(OrderedDict([
        ('_bn2', nn.BatchNorm1d(net._bn1.num_features)),
        ('_drop1', nn.Dropout(p=net._global_params.dropout_rate)),
        ('_linear1', nn.Linear(net._bn1.num_features, 512)),
        ('_relu', nn.ReLU()),
        ('_bn3', nn.BatchNorm1d(512)),
        ('_drop2', nn.Dropout(p=net._global_params.dropout_rate / 2))
    ]))
    fc = nn.Linear(512, net._global_params.num_classes)

    net._dropout = dropout
    net._fc = fc

    data = torch.zeros((2, 3, img_size, img_size))
    output = net(data)
    assert not torch.isnan(output).any()
Test ability to modify dropout and fc modules of network
tests/test_model.py