fixed types

model/train.py

@@ -18,7 +18,7 @@ from torch.optim.lr_scheduler import ExponentialLR
 from torch.utils.data import DataLoader
 
 from model.src.utils import VCDADataset, get_class
-from shared.datastore import connect_minio
+from shared.datastore import Datastore
 
 
 def parse_arguments():
@@ -57,18 +57,19 @@ optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
 loss_fn = nn.CrossEntropyLoss()
 
 # create datasets and loaders
-minio_client = connect_minio()
+datastore = Datastore()
+datastore.connect()
 with open('model/src/dataset/train.csv', encoding='utf-8') as fh:
     train_data_name_list = fh.read().split('\n')
 train_dataset = VCDADataset(
-    minio_client=minio_client,
+    datastore=datastore,
     data_name_list=train_data_name_list,
 )
 train_loader = DataLoader(dataset=train_dataset, num_workers=args.loader_workers)
 
 with open('model/src/dataset/val.csv', encoding='utf-8') as fh:
     val_data_name_list = fh.read().split('\n')
-val_dataset = VCDADataset(minio_client=minio_client, data_name_list=val_data_name_list)
+val_dataset = VCDADataset(datastore=datastore, data_name_list=val_data_name_list)
 val_loader = DataLoader(dataset=val_dataset, num_workers=args.loader_workers)
 
 # create trainer and evaluator