Also, I wanted to check if my model may face vanishing gradient issue, therefore I wanted to plot grad thorough out, my model. Please can you help me how can I access them?
Also, any suggestion on the points on which I can look is strongly welcome.
My model,
class basic_CNN2(nn.Module):
    """
    Basic CNN architecture deployed
    """
    def __init__(self):
        super(basic_CNN2, self).__init__()
        self.Block1 = nn.Sequential(
            nn.Conv2d(3,16,3),
            nn.ReLU(),
            nn.Conv2d(16,16,3),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        self.Block2 = nn.Sequential(
            nn.Conv2d(16,32,3),
            nn.ReLU(),
            nn.Conv2d(32,32,3),
            nn.ReLU(),
            nn.BatchNorm2d(32),
            nn.MaxPool2d(2)
        )
        self.Block3 = nn.Sequential(
            nn.Conv2d(32,64,3),
            nn.ReLU(),
            nn.Conv2d(64,64,3),
            nn.ReLU(),
            nn.BatchNorm2d(64),
            nn.MaxPool2d(2)
        )
        self.Block4 = nn.Sequential(
            nn.Conv2d(64,128,3),
            nn.ReLU(),
            nn.Conv2d(128,128,3),
            nn.ReLU(),
            nn.BatchNorm2d(128),
            nn.MaxPool2d(2),
            nn.Dropout(0.2)
        )
        self.Block5 = nn.Sequential(
            nn.Conv2d(128,256,3),
            nn.ReLU(),
            nn.Conv2d(256,256,3),
            nn.ReLU(),
            nn.BatchNorm2d(256),
            nn.MaxPool2d(2),
            nn.Dropout(0.2)
        )
        self.Linear = nn.Sequential(
            nn.Linear(2304,1024),
            nn.ReLU(),
            nn.Dropout(0.7),
            nn.Linear(1024,512),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(512,64),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(64,1),
            nn.Sigmoid()
        )
        

    def forward(self, x):
        x = self.Block1(x)
        x = self.Block2(x)
        x = self.Block3(x)
        x = self.Block4(x)
        x = self.Block5(x)
        x = x.view(-1,self.num_flat_features(x))
        x = self.Linear(x)
        return x

    def num_flat_features(self, x):
        size = x.size()[1:]
        num_features = 1
        for s in size:
            num_features *= s
        return num_features

Thanks in advance