Update train.py

Author: guillemsimeon

torchmdnet/scripts/train.py

@@ -103,6 +103,7 @@ def get_argparse():
         `a[1] = a[2] = b[2] = 0`;`a[0] >= 2*cutoff, b[1] >= 2*cutoff, c[2] >= 2*cutoff`;`a[0] >= 2*b[0]`;`a[0] >= 2*c[0]`;`b[1] >= 2*c[1]`;
         These requirements correspond to a particular rotation of the system and reduced form of the vectors, as well as the requirement that the cutoff be no larger than half the box width.
     Example: [[1,0,0],[0,1,0],[0,0,1]]""")
+    parser.add_argument('--vector-output', type=bool, default=False, help='If true, returns vector features per atom on top of scalars')
     parser.add_argument('--static_shapes', type=bool, default=False, help='If true, TensorNet will use statically shaped tensors for the network, making it capturable into a CUDA graphs. In some situations static shapes can lead to a speedup, but it increases memory usage.')
 
     # other args