Update crnn_main_v2.py

crnn_main_v2.py

@@ -37,7 +37,7 @@ def weights_init(m):
 
 
 def val(net, val_loader, criterion, iteration, max_i=1000):
-    
+
     print('Start val')
     for p in crnn.parameters():
         p.requires_grad = False
@@ -123,10 +123,14 @@ def main(crnn, train_loader, val_loader, criterion, optimizer):
         print("is best accuracy: {0}".format(accuracy > params.best_accuracy))
         Iteration+=1
 
+def backward_hook(self, grad_input, grad_output):
+    for g in grad_input:
+        g[g != g] = 0   # replace all nan/inf in gradients to zero
+
 if __name__ == '__main__':
 
     # args = init_args()
-    # manualSeed = random.randint(1, 10000)  # fix seed
+    # manualSeed = random.randint(1, 10000)  #fix seed
     manualSeed=10
     random.seed(manualSeed)
     np.random.seed(manualSeed)
@@ -175,4 +179,5 @@ def main(crnn, train_loader, val_loader, criterion, optimizer):
     else:
         optimizer = optim.RMSprop(crnn.parameters(), lr=params.lr)
 
+    crnn.register_backward_hook(backward_hook)
     main(crnn, train_loader, val_loader, criterion, optimizer)