Following this instruction :
1. Strictly adhere to PEP-8 standards, ensuring code is clean, well-structured, and follows Python's style guide for readability.
2. Use appropriate modules like `typing` for type hints, ensuring code clarity, maintainability, and leveraging Python's standard libraries effectively.
3. Focus on time, space, robustness, scalability; write optimized, efficient code that can handle large datasets and scale seamlessly.
4. Implement comprehensive error handling, catching exceptions gracefully, providing meaningful error messages, and ensuring code robustness under unexpected conditions.
5. Ensure the developer agent writes advanced, maintainable Python code, balancing performance and readability while adhering to best practices.



TASK :make below pipeline generalised, training,evaluation,inference,plots and weights save in the ouput folder 

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets,transforms
import timm

transform=transforms.Compose(
[
transforms.RandomHorizontalFlip(),
transforms.RandomVerticalFlip(),
transformers.ToTensor(),
transformers.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))

]
)
train_dataset=datasets.CIFAR10(root="./data",train=True,download=True,transform=transform)
test_dataset=datasets.CIFAR10(root="./data",train=True,download=True,transform=transform)
train_loader=DataLoader(train_dataset,batch_size=64,shuffle=True,num_workers=4)
test_loader=DataLoader(test_dataset,batch_size=64,shuffle=True,num_workers=4)

model=timm.create_model("resnet18",pretrained=True)
model.fc=nn.Linear(model.fc.in_features,10)
criterion=nn.CrossEntropyLoss()
optimizer=optim.SGD(model.parameters(),lr=0.01,momentum=0.9)

device=torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
num_epochs=30
for epoch in range(num_epochs):
    model.train()
    for inputs, labels in train_loader:
        inputs,labels=inputs.to(device),labels.to(device)
        optimizer.zero_grad()
        outputs=model(inputs)
        loss=criterion(outputs,labels)
        loss.backward()
        optimizer.step()
    #Validation
    model.train()
    correct=0
    total=0
    with torch.no_grad():
         for inputs,labels in test_loader:
             inputs,labels =inputs.to(device), labels.to(device)
             outputs=model(inputs)
             _,predicted=torch.max(outputs.data,1)
             total+=labels.size(0)
             correct+=(predicted==labels).sum().item()
    accuracy=correct/total
    print(f"Epoch {epoch+1}/{num_epochs},Loss:{loss:4.4f},Accuracy:(accuracy:.4f})

torch.save(model.state_dict(),"resnet18_cifar10.pth")
    
from PIL import Image
import torchvision.transforms as transforms
model=model.create_model("resnet18")
model.fc=nn.Linear(model.fc.in_features,10)
model.load_state_dict(torch.load('resnet18_cifar10.pth'))
model.eval()

image_transform=transforms.Compose(
[
transformers.Resize((32,32)),
transformers.ToTensor(),
transformers.Normalizr((0.5,0.0,0.5),(0.5,0.5,0.5))


]
)

!wget https://madbarn.com/wp-content/uploads/2023/12/List-of-Gaited-Horse-Breeds.jpg

import requests
from io import BytesIO
import matplotlib.pyplot as plt

Url=""
response=requests.get(Url)
image=Image.open(r"")
plt.imshow(image)

input_image=image_transform(image).unsqueeze(0)
import matplotlib.pyplot as plt
class_names=[
"airplane","automodile",

]
plt.imshow(image)
with torch.no_grad():
     model_ouput=model(input_image)
     _,predicted_class=torch.max(model_output,1)
predicted_class_index=predicted_class.item()
predicted_class-name=