What do I do wrong with validation?

I make model training (which goes fine, with no error, loss is decreasing and accuracy is increasing), however if I add validation part, the cell gives me an error.
index out of range in self
If I try to repeat the error by going through the model code separately with X_val and y_val data – everything works. If I try to run the model as function, I get the same error.
The shape of batches and arrays are the same as for X_train and y_train.

train_batches = round(len(X_train)/batch_size)

for epoch in range(10):

    for train_bat in range (train_batches):
        batch_ix = np.random.randint(len(X_train), size=batch_size)
        batch_inp = X_train.iloc[batch_ix]
        batch_out = torch.tensor(y_train.iloc[batch_ix].values).to(torch.float32)
        cat_tensor = torch.tensor(batch_inp[cat_cols].values).to(torch.int64)
        num_tensor = torch.tensor(batch_inp[num_cols].values).to(torch.int64)
        model_pred = model(num_tensor, cat_tensor)

        train_loss = loss_fn(model_pred, batch_out)    
        train_acc_sum += ((model_pred > 0.0) == batch_out).sum()/len(batch_out)
    with torch.no_grad():
        for val_bat in range (val_batches):
             print(ep, val_bat, i)
             val_batch_ix = np.random.randint(len(X_val), size=batch_size)
             val_batch_inp = X_val.iloc[val_batch_ix]
             val_batch_out = torch.tensor(y_val.iloc[val_batch_ix].values).to(torch.float32)
             val_cat_tensor = torch.tensor(val_batch_inp[cat_cols].values).to(torch.int64)
             val_num_tensor = torch.tensor(val_batch_inp[num_cols].values).to(torch.int64)
             val_out = model(val_cat_tensor, val_num_tensor)
             val_loss = loss_fn(val_out, val_batch_out)
             val_loss_sum += val_loss.item()
             val_acc_sum += ((val_out > 0.0) == batch_out).sum()/len(val_batch_out)

    val_avg_loss=np.round(val_loss_sum/len(val_batches),2)#val loss of all batches of one epoch
    train_avg_loss=np.round(train_loss_sum/train_batches,2)# train loss of all batches of one epoch
    train_avg_acc=np.round(train_acc_sum/train_batches,2)#train acc of all batches of one epoch
    val_avg_acc=np.round(val_acc_sum/len(val_batches),2)#val acc of all batches of one epoch

    print('Epoch {}, train acc  {}, train loss {} , val acc is {}, loss is {}, learning rate is {} '.format(epoch, train_avg_acc, train_avg_loss, val_avg_acc, val_avg_loss, opt.param_groups[0]['lr']))

The model itself it:

                 embedding_dim_dict=None, learning_rate=0.01):

        self.cat_cols = cat_cols
        self.num_cols = num_cols
        self.embeddings, total_embedding_dim = self._create_embedding_layers(
            cat_cols, embedding_size_dict, embedding_dim_dict)
        in_features = len(num_cols) + total_embedding_dim
        self.layers = nn.Sequential(
            nn.Linear(in_features, 128),
            nn.Linear(128, 256),
            nn.Linear(256, n_classes)

    def _create_embedding_layers(cat_cols, embedding_size_dict, embedding_dim_dict):
        """construct the embedding layer, 1 per each categorical variable"""
        total_embedding_dim = 0
        embeddings = {}
        for col in cat_cols:
            embedding_size = embedding_size_dict[col]
            embedding_dim = embedding_dim_dict[col]
            total_embedding_dim += embedding_dim
            embeddings[col] = nn.Embedding(embedding_size, embedding_dim)

        return nn.ModuleDict(embeddings), total_embedding_dim

    def forward(self, num_tensor, cat_tensor):
        cat_outputs = []
        for i, col in enumerate(self.cat_cols):
            embedding = self.embeddings[col]
            cat_output = embedding(cat_tensor[:, i])
        cat_outputs = torch.cat(cat_outputs, dim=1)
        all_outputs = torch.cat((num_tensor, cat_outputs), dim=1)
        final_outputs = self.layers(all_outputs).squeeze(dim=1)
        return final_outputs

And the error happens here when val_out is calculated:

     46             embedding = self.embeddings[col]
---> 47             cat_output = embedding(cat_tensor[:, i])
     48             cat_outputs.append(cat_output)

I rechecked the shape of tensors and removed any other iterations using “i” just in case.

I re-wrote the categorical embeddings preparation, so now it works