model = tf.keras.models.Sequential([  
    tf.keras.layers.Dense(64, input_shape = (384,), activation = "relu"),
    tf.keras.layers.Dense(36, activation="softmax", use_bias = False)
                          ])    
        
model.summary()  
Model: "sequential_5"
_________________________________________________________________
 Layer (type)                Output Shape              Param #   
=================================================================
 dense_13 (Dense)            (None, 64)                24640     
                                                                 
 dense_14 (Dense)            (None, 36)                2304 

# base model - it works but gives v low accuracy. need to add more layers.
class LR(torch.nn.Module):
    def __init__(self, n_features):
        super(LR, self).__init__()
        self.lr = torch.nn.Linear(n_features, 36)
        nn.ReLU()
        
    def forward(self, x):
        out = torch.softmax(self.lr(x),dim = 1,dtype=None)
        return out
# parameters
n_features = 384
n_classes = 36
optim = torch.optim.SGD(model.parameters(), lr=0.1)
criterion = torch.nn.CrossEntropyLoss()
accuracy_fn = Accuracy(task="multiclass",  num_classes=36)

class LR(torch.nn.Module):
    def __init__(self, n_features):
        super(LR, self).__init__()
        self.lr = torch.nn.Linear(n_features, 64)        
        self.lr = nn.ReLU()
        self.lr = torch.nn.Linear(64, 36)
        
    def forward(self, x):
        out = torch.softmax(self.lr(x),dim = 1,dtype=None)
        return out
epochs = 15
def train(model, optim, criterion, x, y, epochs=epochs):
    for e in range(1, epochs + 1):
        optim.zero_grad()
        out = model(x)
        loss = criterion(out, y)
        loss.backward()
        optim.step()
        print(f"Loss at epoch {e}: {loss.data}")
    return model
model = LR(n_features)
model = train(model, optim, criterion, X_train, y_train)
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[175], line 18
     12  #       acc = accuracy_fn(out, y)  
     13  #       acc.backward()
     14  #       optim.step()
     15  #       print(f"Acc at epoch {e}: {acc.data}")
     16     return model
---> 18 model = train(model, optim, criterion, X_train, y_train)
Cell In[175], line 6, in train(model, optim, criterion, x, y, epochs)
      4 for e in range(1, epochs + 1):
      5     optim.zero_grad()
----> 6     out = model(x)
      7     loss = criterion(out, y)
      8     loss.backward()
File /usr/local/lib/python3.10/site-packages/torch/nn/modules/module.py:1501, in Module._call_impl(self, *args, **kwargs)
   1496 # If we don't have any hooks, we want to skip the rest of the logic in
   1497 # this function, and just call forward.
   1498 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1499         or _global_backward_pre_hooks or _global_backward_hooks
   1500         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1501     return forward_call(*args, **kwargs)
   1502 # Do not call functions when jit is used
   1503 full_backward_hooks, non_full_backward_hooks = [], []
Cell In[172], line 10, in LR.forward(self, x)
      9 def forward(self, x):
---> 10     out = torch.softmax(self.lr(x),dim = 1,dtype=None)
     11     return out
File /usr/local/lib/python3.10/site-packages/torch/nn/modules/module.py:1501, in Module._call_impl(self, *args, **kwargs)
   1496 # If we don't have any hooks, we want to skip the rest of the logic in
   1497 # this function, and just call forward.
   1498 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1499         or _global_backward_pre_hooks or _global_backward_hooks
   1500         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1501     return forward_call(*args, **kwargs)
   1502 # Do not call functions when jit is used
   1503 full_backward_hooks, non_full_backward_hooks = [], []
File /usr/local/lib/python3.10/site-packages/torch/nn/modules/linear.py:114, in Linear.forward(self, input)
    113 def forward(self, input: Tensor) -> Tensor:
--> 114     return F.linear(input, self.weight, self.bias)
RuntimeError: mat1 and mat2 shapes cannot be multiplied (155648x384 and 64x36)

X_train.shape
torch.Size([155648, 384]) 
y_train.shape 
torch.Size([155648]).

class LR(torch.nn.Module):
    def __init__(self, n_features):
        super(LR, self).__init__()
        self.lr = torch.nn.Linear(n_features, 64)        
        self.lr = nn.ReLU()                  # lr is over-written here
        self.lr = torch.nn.Linear(64, 36)    # and here

class LR(torch.nn.Module):
    def __init__(self, n_features):
        super(LR, self).__init__()
        self.lr = torch.nn.Linear(64, 36)

from torch import nn
model = nn.Sequential(
    nn.Linear(384, 64),
    nn.ReLU()
    nn.Linear(64, 36, bias=False)
    nn.Softmax()
)

import torch.nn as nn
class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.L1 = nn.Linear(384, 64),
        self.RLU = nn.ReLU()
        self.L2 = nn.Linear(64, 36, bias=False)
    def forward(self, x):
        x = self.L1(x)
        x = self.RLU(x)
        x = self.L2(x)
        return x

class LR(torch.nn.Module):
    def __init__(self, n_features):
        super(LR, self).__init__()
        self.lr = torch.nn.Sequential(
            torch.nn.Linear(n_features, 64),        
            torch.nn.ReLU(),
            torch.nn.Linear(64, 36, bias=False),
            torch.nn.Softmax(dim=1),
        )
    def forward(self, x):
        return self.lr(x)

x = torch.ones(1, 1000)
lr = LR(x.shape[1])
lr(x)