import torch
import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=64, kernel_size=(3, 3), padding=(1, 1), stride=(1, 1))
        self.conv2 = nn.Conv2d(in_channels=64, out_channels=2, kernel_size=(3, 3), padding=(1, 1), stride=(1, 1))
        
        self.linear = nn.Linear(in_features=2*32*32,out_features=2)
    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = x.view(-1)
        x = self.linear(x)
        return x
        
model = Model()
data = torch.randn((3,32,32))
y_pred = model(data) 
y_pred
    

 
tensor([ 0.0984, -0.0103], grad_fn=ViewBackward0)