Update export.py with --train mode argument

[glenn-jocher]

This will allow error-free export of YOLOv5 to CoreML models. All ops are included except for the Detect layer grid reconstruction, which is not used during training. The model will output the 3 or 4 output layers rather than the concatenated reconstructed grid output used during .eval() mode.

🛠️ PR Summary

_{Made with ❤️ by Ultralytics Actions}

🌟 Summary

Enhanced export script with training mode option 🛠️

📊 Key Changes

• Added an argument (--train) to the export script allowing users to specify training mode during export.
• Ensured compatibility with PyTorch 1.6.0 by setting non-persistent buffer sets.

🎯 Purpose & Impact

• Purpose: The addition of --train allows the model to be exported in training mode, which is useful for continuing training or fine-tuning on other datasets.
• Impact: Users will have more flexibility in how they export models, potentially improving the model's adaptability to various tasks. It also assures better forward compatibility with older versions of PyTorch.

[glenn-jocher]

Should resolve #2961