fix my code

from google.colab import drive
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms, datasets
import os

# Assuming Google Drive is already mounted
# If not, uncomment the next line and execute it.
# drive.mount('/content/drive')

# Device configuration
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f'Using {device} device')

class HandSignsDataset(Dataset):
    def __init__(self, root_dir, transform=None):
        self.transform = transform
        self.images = datasets.ImageFolder(root=root_dir, transform=transform)
        self.classes = self.images.classes
        self.num_classes = len(self.classes)

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

    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()
        img, label = self.images[idx]
        return img, label

transform = transforms.Compose([
    transforms.Resize((224, 224)),  # Resize images to a fixed size
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# Specify your paths here
train_data_path = '/content/drive/MyDrive/Output_ISL/train'
test_data_path = '/content/drive/MyDrive/Output_ISL/test'

# Loading the datasets
train_dataset = HandSignsDataset(root_dir=train_data_path, transform=transform)
test_dataset = HandSignsDataset(root_dir=test_data_path, transform=transform)

train_loader = DataLoader(train_dataset, batch_size=10, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=10, shuffle=False)

# Determine the number of classes dynamically
num_classes = train_dataset.num_classes

# Defining the Neural Network
class Net(nn.Module):
    def __init__(self, num_classes):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
        self.fc1 = nn.Linear(64 * 56 * 56, 600)  # Adjusted for the size after pooling
        self.fc2 = nn.Linear(600, num_classes)

    def forward(self, x):
        x = self.pool(nn.functional.relu(self.conv1(x)))
        x = self.pool(nn.functional.relu(self.conv2(x)))
        x = x.view(-1, 64 * 56 * 56)  # Adjusted for the size after pooling
        x = nn.functional.relu(self.fc1(x))
        x = self.fc2(x)
        return x

model = Net(num_classes).to(device)

# Loss and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# Training the model
def train_model(model, criterion, optimizer, num_epochs=10):  # Adjust num_epochs as required
    model.train()
    for epoch in range(num_epochs):
        for i, (images, labels) in enumerate(train_loader):
            images, labels = images.to(device), labels.to(device)

            # Forward pass
            outputs = model(images)
            loss = criterion(outputs, labels)

            # Backward and optimize
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')

# Function to test the model
def test_model(model):
    model.eval()
    with torch.no_grad():
        correct = 0
        total = 0
        for images, labels in test_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

        print(f'Accuracy of the model on the test images: {100 * correct / total}%')

train_model(model, criterion, optimizer)
test_model(model)