diff --git a/distiller/apputils/checkpoint.py b/distiller/apputils/checkpoint.py
index fc90e8877deb6f8da54e491aea29f5e342fca328..f45026b49c3669352622d750779c9adc99dd4c5d 100755
--- a/distiller/apputils/checkpoint.py
+++ b/distiller/apputils/checkpoint.py
@@ -68,7 +68,7 @@ def save_checkpoint(epoch, arch, model, optimizer=None, scheduler=None,
         checkpoint['dataset'] = model.dataset
         if not arch:
             checkpoint['arch'] = model.arch
-    except NameError:
+    except AttributeError:
         pass
 
     if optimizer is not None: