quick fix for transformer_flops.py error

[jahatef]

Fixes this error when running transformer_flops.py. the train_batch_size wasn't being set, causing a weird error.

[2024-02-16 02:42:42,879] [INFO] [real_accelerator.py:161:get_accelerator] Setting ds_accelerator to cuda (auto detect)
Traceback (most recent call last):
  File "/weka/home-jacob/cookbook/sizing/benchmarks/transformer_flops.py", line 276, in <module>
    megatron_wrapper.initialize_megatron(configurations[0])
  File "/weka/home-jacob/cookbook/sizing/benchmarks/megatron_wrapper.py", line 92, in initialize_megatron
    args = get_megatron_args(configuration, override_tensor_mp_size=True)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/weka/home-jacob/cookbook/sizing/benchmarks/megatron_wrapper.py", line 114, in get_megatron_args
    neox_args = megatron.NeoXArgs.from_dict(asdict(args))
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/weka/home-jacob/cookbook/sizing/benchmarks/megatron/neox_arguments/arguments.py", line 234, in from_dict
    return cls(**args_dict)
           ^^^^^^^^^^^^^^^^
  File "<string>", line 228, in __init__
  File "/weka/home-jacob/cookbook/sizing/benchmarks/megatron/neox_arguments/arguments.py", line 134, in __post_init__
    if not self.validate_types():
           ^^^^^^^^^^^^^^^^^^^^^
  File "/weka/home-jacob/cookbook/sizing/benchmarks/megatron/neox_arguments/arguments.py", line 1238, in validate_types
    elif field_def.type.__origin__ == Literal:
         ^^^^^^^^^^^^^^^^^^^^^^^^^
AttributeError: type object 'int' has no attribute '__origin__'