Update train.py

awslabs · Jan 14, 2025 · e794b8e · e794b8e
1 parent f4d3599
commit e794b8e
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/tests/end2end-tests/custom-gnn/train.py b/tests/end2end-tests/custom-gnn/train.py
@@ -108,6 +108,7 @@ def main(args):
                            help="The name of the label.")
     argparser.add_argument("--num-classes", type=int, required=True,
                            help="The number of classes.")
+    # compatible for all torch versions
     argparser.add_argument("--local-rank", "--local_rank", type=int,
                            help="The rank of the trainer.")
     argparser.add_argument("--verbose",