From 7787090165e5473ab8dae2a89e2ac51178c4bae3 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Wed, 17 Apr 2019 18:33:16 +0200
Subject: [PATCH] updates

---
 train.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/train.py b/train.py
index 13268620..bbfb55f5 100644
--- a/train.py
+++ b/train.py
@@ -113,7 +113,7 @@ def train(
     # Dataloader
     dataloader = DataLoader(dataset,
                             batch_size=batch_size,
-                            num_workers=opt.num_workers,
+                            num_workers=0,
                             shuffle=True,
                             pin_memory=True,
                             collate_fn=dataset.collate_fn,
@@ -170,6 +170,9 @@ def train(
 
             # Compute loss
             loss, loss_items = compute_loss(pred, targets, model)
+            if torch.isnan(loss):
+                print('WARNING: nan loss detected, ending training')
+                return results
 
             # Compute gradient
             if mixed_precision: