Tensorboard model visualization bug fix
diff --git a/train.py b/train.py
index 1f2b467e732..82043b7fff3 100644
--- a/train.py
+++ b/train.py
@@ -332,7 +332,7 @@ def train(hyp, opt, device, tb_writer=None):
                     Thread(target=plot_images, args=(imgs, targets, paths, f), daemon=True).start()
                     # if tb_writer:
                     #     tb_writer.add_image(f, result, dataformats='HWC', global_step=epoch)
-                    #     tb_writer.add_graph(model, imgs)  # add model to tensorboard
+                    #     tb_writer.add_graph(torch.jit.trace(model, imgs, strict=False), [])  # add model graph
                 elif plots and ni == 10 and wandb_logger.wandb:
                     wandb_logger.log({"Mosaics": [wandb_logger.wandb.Image(str(x), caption=x.name) for x in
                                                   save_dir.glob('train*.jpg') if x.exists()]})

This fix should allow for visualizing YOLOv5 model graphs correctly in Tensorboard by uncommenting line 335 in train.py:
```python
if tb_writer:
    tb_writer.add_graph(torch.jit.trace(model, imgs, strict=False), [])  # add model graph
```

The problem identified in #2284 was that the detect() layer checks the input size to adapt the grid if required, and tracing does not seem to like this shape check (even if the shape is fine and no grid recomputation is required). The following will warn:
https://github.com/ultralytics/yolov5/blob/0cae7576a9241110157cd154fc2237e703c2719e/train.py#L335

Solution is below. This is a YOLOv5s model displayed in TensorBoard. You can see the Detect() layer merging the 3 layers into a single output for example, and everything appears to work and visualize correctly.
```python
tb_writer.add_graph(torch.jit.trace(model, imgs, strict=False), [])
```
<img width="893" alt="Screenshot 2021-04-11 at 01 10 09" src="https://user-images.githubusercontent.com/26833433/114286928-349bd600-9a63-11eb-941f-7139ee6cd602.png">

## 🛠️ PR Summary

<sub>Made with ❤️ by [Ultralytics Actions](https://github.com/ultralytics/actions)<sub>

### 🌟 Summary
Optimized TensorBoard model graph logging with JIT tracing.

### 📊 Key Changes
- Commented-out code for adding the unoptimized model graph to TensorBoard has been replaced with JIT-traced model graph.
- This change is in the `train.py` file, specifically in the model training loop.

### 🎯 Purpose & Impact
- 🚀 **Purpose**: To enable more efficient and possibly more accurate graph logging in TensorBoard.
- 📈 **Impact**: Users gain the ability to visualize a traced version of their model's graph, which might be faster and more optimized than directly logging the model.
- 🧑‍💻 **For Developers**: This could help with better understanding model performance and optimizations applied by JIT.
- 🔄 **For Users**: Insights from detailed visualizations might help in tweaking the model and potentially achieving better results.