Review that code and tell me if anything seems wrong:

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import numpy as np

# Function to set random seed for reproducibility
def set_random_seed(seed):
    if seed is None:
        return None
    torch.manual_seed(seed)
    np.random.seed(seed)

# Custom Dataset to handle the data structure
class CustomDataset(Dataset):
    def __init__(self, x, y, sample_weight):
        self.x = x
        self.y = y
        self.sample_weight = sample_weight

    def __len__(self):
        return len(self.y)

    def __getitem__(self, idx):
        return self.x[idx], self.y[idx], self.sample_weight[idx]

# Logistic Regression Model
class LogisticRegressionModel(nn.Module):
    def __init__(self, input_dim, output_dim):
        super(LogisticRegressionModel, self).__init__()
        self.linear = nn.Linear(input_dim, output_dim)

    def forward(self, x):
        return self.linear(x)

# Main Logistic Regression class
class LogisticRegression:
    def __init__(self, random_seed=42, num_epochs=100, learning_rate=0.01, batch_size=32, patience=5, min_delta=0.0):
        self.model = None
        self.random_seed = random_seed
        self.data = {}
        self.num_epochs = num_epochs
        self.learning_rate = learning_rate
        self.batch_size = batch_size
        self.patience = patience
        self.min_delta = min_delta

    def fit(self, X, y, sample_weight=None):
        # Set random seed for reproducibility
        if self.random_seed is not None:
            set_random_seed(self.random_seed)

        # Convert input data to tensors
        X = torch.tensor(X).float()
        y = torch.tensor(y).long()
        if sample_weight is not None:
            sample_weight = torch.tensor(sample_weight).float()
        else:
            sample_weight = torch.ones(len(y)).float()

        # Determine input and output dimensions
        input_dim = X.shape[1]
        output_dim = torch.unique(y).numel()

        # Initialize the model
        self.model = LogisticRegressionModel(input_dim, output_dim)

        # Sample some of the good training data for early stopping validation and weights_val
        good_indices = np.where(sample_weight == 1)[0]
        bad_indices = np.where(sample_weight == 2)[0]
        num_good_samples = len(good_indices)
        num_val_samples = max(1, int(0.1 * num_good_samples))
        num_good_train_samples = num_good_samples - num_val_samples
        num_weights_val_samples = max(1, int(0.2 * num_good_train_samples))

        # Ensure reproducibility in sampling
        indices = torch.randperm(num_good_samples, generator=torch.Generator().manual_seed(self.random_seed))
        train_indices = np.concatenate((good_indices[indices[:num_good_train_samples].tolist()], bad_indices))
        val_indices = good_indices[indices[num_good_train_samples:]]
        weights_val_indices = good_indices[indices[:num_weights_val_samples]]

        self.data['train'] = {
            'x': X[train_indices],
            'y': y[train_indices],
            'sample_weight': sample_weight[train_indices]
        }
        
        self.data['early_stopping_validation'] = {
            'x': X[val_indices],
            'y': y[val_indices],
            'sample_weight': sample_weight[val_indices]
        }
        
        self.data['weights_val_validation'] = {
            'x': X[weights_val_indices],
            'y': y[weights_val_indices],
            'sample_weight': sample_weight[weights_val_indices]
        }

        # Create datasets and data loaders
        train_dataset = CustomDataset(
            self.data['train']['x'],
            self.data['train']['y'],
            self.data['train']['sample_weight']
        )
        val_dataset = CustomDataset(
            self.data['early_stopping_validation']['x'],
            self.data['early_stopping_validation']['y'],
            self.data['early_stopping_validation']['sample_weight']
        )
        weights_val_dataset = CustomDataset(
            self.data['weights_val_validation']['x'],
            self.data['weights_val_validation']['y'],
            self.data['weights_val_validation']['sample_weight']
        )

        train_loader = DataLoader(train_dataset, batch_size=self.batch_size, shuffle=True)
        val_loader = DataLoader(val_dataset, batch_size=self.batch_size, shuffle=False)
        weights_val_loader = DataLoader(weights_val_dataset, batch_size=len(weights_val_dataset), shuffle=False)#! one update
        print(f"Training on {len(train_dataset)} samples, validating on {len(val_dataset)} samples, and weights validating is {len(weights_val_dataset)} samples.")
        
        # Define loss function and optimizer
        criterion = nn.CrossEntropyLoss(reduction='none')  # 'none' to apply sample weights manually
        optimizer = optim.SGD(self.model.parameters(), lr=self.learning_rate)

        # Variables for early stopping
        best_val_loss = float('inf')
        counter = 0
        def forward(model, X_batch, y_batch, weights_array, loss_fn):
            outputs = model(X_batch)
            losses = loss_fn(outputs, y_batch)
            if weights_array == None:
                return losses
            return losses * weights_array

        weights_per_epoch = []
        # Training loop
        for epoch in range(self.num_epochs):
            # Training phase
            self.model.train()
            total_train_loss = 0.0
            for xb, yb, wb in train_loader:
                εs = torch.zeros(2, requires_grad=True)
                weights_array = wb
                indices_list = []
                for i in range(len(εs)):
                    indices = (weights_array == (i + 1)).nonzero(as_tuple=True) # Get the indices for the different shades
                    indices_list.append(indices) # Store these indices for later use
                    weights_array[indices] = εs[i]
                optimizer.zero_grad()
                with higher.innerloop_ctx(self.model, optimizer) as (meta_model, meta_opt):
                    meta_loss = forward(meta_model, xb, yb, weights_array, criterion)
                    meta_loss = meta_loss.mean()
                    meta_opt.step(meta_loss)
                    total_weights_eval_loss = torch.tensor(0.0)
                    num_weights_eval_loss=0
                    for xb_w, yb_w, wb_w in weights_val_loader:
                        losses = forward(meta_model, xb_w, yb_w, None, criterion)
                        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 = torch.zeros(len(wb))
                for i in range(len(εs)):
                    weights_array[indices_list[i]] = weights[i]
                weighted_loss = forward(self.model, xb, yb, weights_array, criterion)
                weighted_loss = weighted_loss.mean()
                total_train_loss += weighted_loss.item()
                weighted_loss.backward()
                optimizer.step()
                weights_per_epoch.append(weights.cpu().detach().numpy())

            avg_train_loss = total_train_loss / len(train_dataset)

            # Validation phase
            self.model.eval()
            total_val_loss = 0.0
            with torch.no_grad():
                for xb, yb, wb in val_loader:
                    outputs = self.model(xb)
                    loss = criterion(outputs, yb)
                    loss = loss.mean()
                    total_val_loss += loss.item() * xb.size(0)

            avg_val_loss = total_val_loss / num_val_samples

            # Early stopping check
            # If the loss improved more than min_delta, reset the counter
            if best_val_loss - avg_val_loss > self.min_delta:
                best_val_loss = avg_val_loss
                counter = 0
                best_model_state = self.model.state_dict()
            else:
                counter += 1
                if counter >= self.patience:
                    print(f"Early stopping at epoch {epoch + 1}")
                    break
            print(np.average(weights_per_epoch, axis=0))
            print(f"Epoch {epoch + 1}/{self.num_epochs} - "
                  f"Train Loss: {avg_train_loss:.4f} - Val Loss: {avg_val_loss:.4f}")

        # Load the best model state
        self.model.load_state_dict(best_model_state)

    def predict(self, X):
        self.model.eval()
        X = torch.tensor(X).float()
        with torch.no_grad():
            outputs = self.model(X)
            _, predicted = torch.max(outputs, 1)
        return predicted.numpy()

    def predict_proba(self, X):
        self.model.eval()
        X = torch.tensor(X).float()
        with torch.no_grad():
            outputs = self.model(X)
            probabilities = nn.functional.softmax(outputs, dim=1)
        return probabilities.numpy()