devanshuDesai

import torch
from torch import nn
from tqdm import tqdm_notebook

class CNN(nn.Module):
    def __init__(self, input_size, num_classes):
        """
        init convolution and activation layers
        Args:
            input_size: (1,28,28)