From 1cfd28074d1a5bee07a95b96d31d31b0152d893c Mon Sep 17 00:00:00 2001
From: Nina1031 <73870119+Nina1031@users.noreply.github.com>
Date: Tue, 25 May 2021 16:53:13 +0800
Subject: [PATCH] Update train.py
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

感觉 dataset_train 的cfg里选择了训练类别
---
 train.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/train.py b/train.py
index cfe3b3ee..a436c2d4 100644
--- a/train.py
+++ b/train.py
@@ -161,7 +161,7 @@ def main(cfg, gpus):
             net_encoder, net_decoder, crit)
 
     # Dataset and Loader
-    dataset_train = TrainDataset(
+    dataset_train = TrainDataset(       # select trian index
         cfg.DATASET.root_dataset,
         cfg.DATASET.list_train,
         cfg.DATASET,