Update train.py

train.py

@@ -59,7 +59,7 @@ print(f"using {device} device")
 
 
 
-class TensorMapperImageClassification(LiteTensorMapper):
+class LiteTensorMapperImageClassification(LiteTensorMapper):
     def __init__(
         self,
         image_size=32,
@@ -89,7 +89,7 @@ class TensorMapperImageClassification(LiteTensorMapper):
         out = self.classifier(embedding)
         return out
 
-model = TensorMapperImageClassification().to(device)
+model = LiteTensorMapperImageClassification().to(device)
 print(model)