Update train.py

fffffgggg54 · Jul 31, 2024 · faba057 · faba057
1 parent fdaeca0
commit faba057
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/train.py b/train.py
@@ -485,7 +485,7 @@ def main():
  bn_eps=args.bn_eps,
  scriptable=args.torchscript,
  checkpoint_path=args.initial_checkpoint,
- features_only=args.use_pyramid_head,
+ features_only=args.use_pyramid_head or args.use_fp2t,
  **factory_kwargs,
  **args.model_kwargs,
  )