Remove unused lines

source/model.py

@@ -1,7 +1,5 @@
 import torch.nn as nn
 import torch.nn.functional as F
-from dataset import get_paths, get_data_loader, Dataset
-from setup import Setup
 
 
 class CNN(nn.Module):
@@ -75,27 +73,3 @@ class CNN(nn.Module):
         # print('Out: ', x.size())
         return F.log_softmax(x, dim=1)
 
-if __name__ == '__main__':
-    """
-    Main script to initialize the setup, load datasets, create DataLoader,
-    instantiate the CNN model, and display the number of trainable parameters 
-    and the output size for a batch of images.
-    """
-
-    setup = Setup()
-
-    normal_train_paths, red_train_paths, normal_test_paths, red_test_paths = get_paths()
-
-    train_dataset = Dataset(red_train_paths, normal_train_paths)
-    train_loader = get_data_loader(train_dataset, batch_size=setup.BATCH)
-
-    imgs, labels = next(iter(train_loader))
-
-    cnn = CNN()
-    print(f'Number of trainable parameters in CNN: {sum(p.numel() for p in cnn.parameters() if p.requires_grad)}')
-    output = cnn.forward(imgs)
-
-    # Print info
-    print('\nBatch size: ', setup.BATCH)
-    print('Images size: ', imgs.size())         # (batch, 3, 32, 32)
-    print('CNN output size: ', output.size())   # (batch, 2)