Update run_nn.py

fixed multi-gpu
aiedward · Mar 10, 2018 · 6992ecd · 6992ecd
1 parent ec534da
commit 6992ecd
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/run_nn.py b/run_nn.py
@@ -201,6 +201,10 @@
 
 
    [loss,err,pout] = net(inp,lab,test_flag)
+
+   if multi_gpu:
+     loss=loss.mean()
+     err=err.mean()
 
    if do_forward:
     if rnn==1: