main_densenet.py

# a block consists of initial conv layers followed by 6 dense layers
dense_block = nn.Sequential(
    nn.Conv2d(3, 64, kernel_size=7, padding=3, stride=2, bias=False),
    nn.BatchNorm2d(64),
    nn.MaxPool2d(3, 2),
    Dense_Block(6, 64, growthrate=32, bn_size=4),
)

inputs = torch.rand(1, 3, 100, 100)
outputs = dense_block(inputs)
print(outputs.shape)    # shape would be (1, 256, 24, 24)

# one could also use pretrained weights of DenseNet trained on ImageNet
densenet121 = torchvision.models.densenet121(pretrained=True)