Thank you @m3tobom_M . After flattening, I get this tensor size - torch.Size([64, 19200])
Is that ok? Can I change it to be 1, 19200?
Secondly, I also get this weird error
RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn

For reference, this is my current code:
import torch
from torch import tensor
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader

import numpy as np
from PIL import Image
import time

start = time.time()
transform = transforms.Compose([transforms.Resize((200,200)),transforms.CenterCrop(80),transforms.ToTensor()]) 
training = torchvision.datasets.ImageFolder(root = "training_set", transform = transform)
print(training)

train_dataloader = DataLoader(dataset = training, shuffle = True, batch_size = 64)
print(train_dataloader)




print("Before NN creation",start-time.time())
# Creating the model
class NeuralNetwork (nn.Module):
    def __init__(self,input_dims):
        super(NeuralNetwork, self).__init__()
        self.Linear_Stack = nn.Sequential(
            nn.Linear(input_dims, 1000), 
            nn.ReLU(),
            nn.Linear(1000, 5),
            nn.ReLU(),
            nn.Linear(5, 1),
            nn.LogSoftmax(dim=1)
        )
    def forward (self, x):
        logits = self.Linear_Stack(x)
        return logits

model = NeuralNetwork(19200)
print(model)


epochs = 1

lr = 0.001
loss = nn.CrossEntropyLoss()
optim = torch.optim.SGD(model.parameters(), lr = lr)
print(loss, optim)  

out = 0

def process(tensorin):
    global out
    _, c,h,w = tensorin.shape
    first_layer = c * h * w
    out = tensorin.view(-1, first_layer)
    return out
#creating the main loop
def main(loader, model, loss, optim):
    size = len(loader.dataset)
    print(size)
    for batch_idx, (data, target) in enumerate(loader):
        print(data.shape)
        process(data)
        print(target)
        print(out.shape)
        pred = model(out)
        loss = loss(out, target)
        #backprop
        optim.zero_grad()
        loss.backward()
        optim.step()
        print(loss)

for x in range(epochs):
    main(train_dataloader, model, loss, optim)
    print(start-time.time())

Thank you for any help (and help given - @m3tobom_M , @eqy, @ptrblck  and @ptrmcl)