diff --git a/CNN/dataset.py b/CNN/dataset.py
new file mode 100644
index 0000000..1d6edd7
--- /dev/null
+++ b/CNN/dataset.py
@@ -0,0 +1,47 @@
+import pandas as pd 
+from torch.utils.data import Dataset
+import torch 
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from torch.utils.data import DataLoader
+from torch.utils.data import Dataset
+import torchvision.transforms as transforms
+import torchvision
+
+class TinyImageNet(Dataset):
+    def __init__(self, path, shape , transform=None, target_transform=None):
+        self.imgs , self.img_labels = self.read_data(path,shape)
+        self.transform = transform
+        self.target_transform = target_transform
+        self.num_classes = 20
+
+    def __len__(self):
+        return len(self.imgs)
+
+    def convert_to_onehot(self,indices, num_classes):
+        output = np.eye(num_classes)[np.array(indices).reshape(-1)]
+        result = output.reshape(list(np.shape(indices))+[num_classes])
+        return result    
+
+    def __getitem__(self, idx):
+        image = self.imgs[idx]
+        image = np.transpose(image , (2,0,1))
+        label = self.img_labels[idx]
+        return  torch.from_numpy(image), label
+
+    def read_data(self,path, shape):
+        h = shape[0]
+        w = shape[1]
+        c = shape[2]
+        data = pd.read_csv(path)
+        data = data.to_numpy()
+        # TODO: (0) Check this normalization and look for a better one
+        X = (data[:, 1:-1])/255
+        # TODO: (0) Check this reshaping 
+        X = X.reshape(-1,h,w,c)
+        Y = data[:,-1]
+        #Y = self.convert_to_onehot(Y, num_classes)
+        print("Shape of data: ", np.shape(X), "Shape of train labels: ", np.shape(Y))
+        # print("Shape of data: ", np.type(X), "Shape of train labels: ", np.type(Y))
+        return   np.float32(X)  ,   Y     
diff --git a/CNN/main.py b/CNN/main.py
new file mode 100644
index 0000000..7ac509b
--- /dev/null
+++ b/CNN/main.py
@@ -0,0 +1,80 @@
+import matplotlib.pyplot as plt 
+import numpy as np 
+import os
+import pandas as pd
+ 
+import sys
+import torch 
+import torch.nn as nn
+
+import torch.optim as optim
+from torch.utils.data import DataLoader
+
+
+
+from dataset import TinyImageNet
+from utils import move_to_gpu , get_arg_parser , set_seed
+from model import CNN_Net 
+
+def train(args):
+    train_dataset = TinyImageNet(path=args.train,shape=(64,64,3))
+    valid_dataset = TinyImageNet(path=args.valid,shape=(64,64,3))       
+
+    train_dataloader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True)
+    test_dataloader = DataLoader(valid_dataset, batch_size=4, shuffle=False)
+
+
+    net = CNN_Net()
+
+    move_to_gpu(net)
+    criterion = nn.CrossEntropyLoss()
+    optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=0.9)
+    for epoch in range(args.batch_size):  # loop over the dataset multiple times
+        running_loss = 0.0
+        for i, data in enumerate(train_dataloader, 0):
+            # get the inputs; data is a list of [inputs, labels]
+            inputs, labels = data
+            inputs, labels = inputs.cuda().to(dtype=torch.float) , labels.cuda().to(dtype=torch.long)
+
+            # zero the parameter gradients
+            optimizer.zero_grad()
+
+            # forward + backward + optimize
+            outputs = net(inputs)
+            loss = criterion(outputs, labels)
+            loss.backward()
+            optimizer.step()
+
+            # print statistics
+            running_loss += loss.item()
+            if i % 2000 == 1999:    # print every 2000 mini-batches
+                print('[%d, %5d] loss: %.3f' %
+                    (epoch + 1, i + 1, running_loss / 2000))
+                running_loss = 0.0
+
+
+def accuracy(data_loader): 
+    with torch.no_grad():
+        correct=0
+        total =0 
+        for data in data_loader:
+            inputs, labels = data
+            if torch.cuda.is_available:
+                inputs, labels = inputs.cuda().to(dtype=torch.float) , labels.cuda().to(dtype=torch.long)
+            
+            outputs = net(inputs)
+            predicitons = torch.argmax(outputs,1)
+            correct = correct + torch.sum(predicitons == labels)
+            total = total + labels.size()[0]
+    return correct/total
+
+
+def main():
+    set_seed()
+    parser = get_arg_parser()
+    args = parser.parse_args()
+    print(args)
+    train(args)
+
+if __name__== "__main__":
+    main()
\ No newline at end of file
diff --git a/CNN/model.py b/CNN/model.py
new file mode 100644
index 0000000..2df04fc
--- /dev/null
+++ b/CNN/model.py
@@ -0,0 +1,34 @@
+import torch 
+import torch.nn as nn
+
+class CNN_Net(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = nn.Conv2d(3, 32, (5,5))
+        self.conv2 = nn.Conv2d(32, 32, (5,5))
+        self.pool1 = nn.MaxPool2d(2, 2)
+        self.conv3 = nn.Conv2d(32, 64, (3,3))
+        self.conv4 = nn.Conv2d(64, 64, (3,3))
+        self.pool2 = nn.MaxPool2d(2, 2)
+        self.conv5 = nn.Conv2d(64, 64, (3,3))
+        self.conv6 = nn.Conv2d(64, 128, (3,3))
+        self.pool3 = nn.MaxPool2d(2, 2)
+        self.fc1 = nn.Linear(128*4*4, 256)
+        # self.fc1 = nn.Linear(6272, 256)
+        self.fc2 = nn.Linear(256, 20)
+
+    def forward(self,x):
+        x = F.relu(self.conv1(x))
+        x = F.relu(self.conv2(x))
+        x= self.pool1(x)
+        x = F.relu(self.conv3(x))
+        x = F.relu(self.conv4(x))
+        x= self.pool2(x)
+        x = F.relu(self.conv5(x))
+        x = F.relu(self.conv6(x))
+        x= self.pool3(x)
+
+        x = torch.flatten(x, 1) 
+        x = F.relu(self.fc1(x))
+        x = F.relu(self.fc2(x))
+        return x
diff --git a/CNN/readme.md b/CNN/readme.md
new file mode 100644
index 0000000..342099f
--- /dev/null
+++ b/CNN/readme.md
@@ -0,0 +1,40 @@
+# CNNAssginment 4
+
+This is a programming assignment in which convolutional neural networks are implemented and are trained in Pytorch. 
+
+* Problem statement
+* Train data
+* Validation data
+* Test data
+
+
+## Contents
+train.py
+
+This script contains code for implementation and training of different convolutional neural networks for classification. The network to be trained must be changed within the code. Adam optimizer is used to train the network with cross entropy as the loss function. Data augmentation uses simple tricks such as flipping the images vertically, horizontally and rotating hte images.
+
+Usage
+Run as
+```
+python train.py --lr <learning_rate>  --batch_size <batch_size> --init <init> --save_dir <path_save_dir> --epochs <num_epochs> --dataAugment <augmentation> --train <path_to_train> --val <path_to_val> --test <path_to_test>
+```
+
+* learning_rate: learning rate to be used for all updates, defaults to 0.001
+* batch_size: size of minibatch, defaults to 256
+* init: initialization, 1 corresponds to Xavier and 2 corresponds to He initialization, defaults to 1
+* path_save_dir: path to the folder where the final model is stored
+* num_epochs: number of epochs to run for, defaults to 10
+* augmentation: set to 0 for no augmentation, 1 for augmentation
+* path_to_train: path to the training data .csv file
+* path_to_val: path to the validation dataset .csv file
+* path_to_test: path to the test dataset .csv file
+
+Outputs
+
+
+Note: ./ indicates that the file is created in the working directory
+
+run.sh
+A shell file containing the best set of hyperparameters for the given task. Run as described below to train a network with the specified architecture and predict values for the test data.
+
+./run.sh
\ No newline at end of file
diff --git a/CNN/utils.py b/CNN/utils.py
new file mode 100644
index 0000000..c6cc676
--- /dev/null
+++ b/CNN/utils.py
@@ -0,0 +1,42 @@
+import argparse
+import torch
+import random
+import numpy as np
+
+def str2bool(v):
+    if isinstance(v, bool):
+        return v
+    if v.lower() in ('yes', 'true', 't', 'y', '1'):
+        return True
+    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
+        return False
+    else:
+        raise argparse.ArgumentTypeError('Boolean value expected.')
+
+def get_arg_parser():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--lr", default = 0.001, help = "learning rate, defaults to 0.01", type = float)
+    parser.add_argument("--batch_size", default = 256, help = "size of each minibatch, defaults to 256", type = int)
+    parser.add_argument("--init", default = 1, help = "initialization to be used; 1: Xavier; 2: He; defaults to 1", type = int)
+    parser.add_argument("--save_dir", help = "location for the storage of the final model")
+    parser.add_argument("--epochs", default = 10, help = "number of epochs", type = int)
+    parser.add_argument("--dataAugment", default = 0, help = "1: use data augmentation, 0: do not use data augmentation", type = int)
+    parser.add_argument("--train", default = "train.csv", help = "path to the training data")
+    parser.add_argument("--val", default = "valid.csv", help = "path to the validation data")
+    parser.add_argument("--test", default = "test.csv", help = "path to the test data")
+    
+    return parser
+
+def set_seed(seed = 31):
+    random.seed(seed)
+    torch.manual_seed(seed)
+    np.random.seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
+
+def move_to_gpu(*args):
+    if torch.cuda.is_available():
+        for item in args:
+            item.cuda()