Update train.py

train.py

@@ -7,7 +7,7 @@ from torch import nn
 from torch.utils.data import DataLoader 
 from torchvision import datasets 
 from torchvision.transforms import ToTensor, Normalize, RandomCrop, RandomHorizontalFlip, Compose 
-from cobra import COBRA
+from mcdpmamba import MCDPMAMBA
 # data transforms
 
 transform = Compose([
@@ -63,7 +63,7 @@ print(f"using {device} device")
 
 # model definition
 
-class COBRAImageClassification(COBRA):
+class MCDPMAMBAImageClassification(MCDPMAMBA):
     def __init__(
         self,
         image_size=32,
@@ -93,7 +93,7 @@ class COBRAImageClassification(COBRA):
         out = self.classifier(embedding)
         return out
 
-model = COBRAImageClassification().to(device)
+model = MCDPMAMBAImageClassification().to(device)
 print(model)
 
 # Optimizer
@@ -163,7 +163,7 @@ def test(dataloader, model, loss_fn):
 
 # apply train and test
 
-logname = "/PATH/COBRA/Experiments_cifar10/logs_cobra/logs_cifar10.csv"
+logname = "/PATH/MCDP_MAMBA/Experiments_cifar10/logs_mcdpmamba/logs_cifar10.csv"
 if not os.path.exists(logname):
   with open(logname, 'w') as logfile:
     logwriter = csv.writer(logfile, delimiter=',')
@@ -185,8 +185,8 @@ print("Done!")
 
 # saving trained model
 
-path = "/PATH/COBRA/Experiments_cifar10/weights_cobra"
-model_name = "COBRAImageClassification_cifar10"
+path = "/PATH/MCDP_MAMBA/Experiments_cifar10/weights_mcdpmamba"
+model_name = "MCDPMAMBAImageClassification_cifar10"
 torch.save(model.state_dict(), f"{path}/{model_name}.pth")
 print(f"Saved Model State to {path}/{model_name}.pth ")