is there anything the could be wrong or suspicious in this implementation:
from tqdm import tqdm

def forward_function(model, batch, weights_array, loss_fn, device):
    batch_size, label_seq_len = batch['labels'].shape
    softmax_dim = model.config.vocab_size
    batch = {k: v.to(device) for k, v in batch.items()}
    outputs = model(input_ids=batch['input_ids'], attention_mask=batch['attention_mask'], decoder_input_ids=batch['decoder_input_ids'])
    preds, labels = outputs.logits, batch['labels']
    flattened_preds = preds.view(-1, softmax_dim)
    flattened_labels = labels.view(-1)
    losses = loss_fn(flattened_preds, flattened_labels)
    if weights_array!=None:
      losses_re = losses.view(batch_size, label_seq_len) * weights_array
    else:
      losses_re = losses.view(batch_size, label_seq_len)
    losses_re = losses_re.view(-1)
    losses = losses_re[flattened_labels != -100]
    return losses


def train_one_epoch(model, train_dataloader, device, optimizer, loss_fn):
    model.train()
    total_loss = 0
    batch_losses = []
    train_progress = tqdm(train_dataloader, desc="Training")
    weights_per_epoch = []
    for train_batch in train_progress:
        batch_size, label_seq_len = train_batch['labels'].shape
        εs = torch.zeros(len(λs), requires_grad=True, device=device)
        weights_array = train_batch['weight'].unsqueeze(1).repeat(1, label_seq_len).to(device)
        indices_list = []
        for i in range(len(εs)):
          indices = (weights_array == (i + 1)).nonzero(as_tuple=True) # Get the indices where the weight is equal to i + 1
          indices_list.append(indices) # Store these indices for later use
          weights_array[indices] = εs[i]
        optimizer.zero_grad()
        with higher.innerloop_ctx(model, optimizer) as (meta_model, meta_opt):
          meta_loss = forward_function(meta_model, train_batch, weights_array, loss_fn, device)
          meta_loss = meta_loss.mean()
          meta_opt.step(meta_loss)
          total_weights_eval_loss = torch.tensor(0.0).to(device)
          num_weights_eval_loss=0
          for weight_eval_batch in weight_eval_dataloader:
            losses = forward_function(meta_model, weight_eval_batch, None, loss_fn, device)
            total_weights_eval_loss+=losses.sum()
            num_weights_eval_loss+=len(losses)
          total_weights_eval_loss = total_weights_eval_loss/num_weights_eval_loss
          εs_grad = torch.autograd.grad(total_weights_eval_loss, εs)[0].detach()
        weights = torch.clamp(-εs_grad, min=0)
        weights_norm = torch.sum(weights)
        weights_norm = weights_norm if weights_norm != 0 else 1
        weights = weights / weights_norm
        weights_array = train_batch['weight'].unsqueeze(1).repeat(1, label_seq_len).to(device)
        for i in range(len(εs)):
          weights_array[indices_list[i]] = weights[i]
        weighted_loss = forward_function(model, train_batch, weights_array, loss_fn, device)
        weighted_loss = weighted_loss.mean()
        total_loss += weighted_loss.item()
        batch_losses.append(weighted_loss.item())
        weighted_loss.backward()
        optimizer.step()
        train_progress.set_postfix(loss=weighted_loss.item())
        weights_per_epoch.append(weights.cpu().detach().numpy())

    avg_train_loss = total_loss / len(train_dataloader)
    return avg_train_loss, batch_losses, np.average(weights_per_epoch, axis=0)