diff --git a/__dev/port.py b/__dev/port.py
index 8d6bc06..50887bf 100644
--- a/__dev/port.py
+++ b/__dev/port.py
@@ -12,7 +12,7 @@ import os
 import time
 import code
 
-class TorchBatchNorm2D(keras.engine.topology.Layer):
+class TorchBatchNorm2D(keras.engine.Layer):
     def __init__(self, axis=-1, momentum=0.99, epsilon=1e-3, **kwargs):
         super(TorchBatchNorm2D, self).__init__(**kwargs)
         self.supports_masking = True
@@ -71,6 +71,7 @@ fa = face_alignment.FaceAlignment(face_alignment.LandmarksType._2D,enable_cuda=F
 fa.eval()
 
 
+
 def KerasConvBlock(in_planes, out_planes, input, srctorch):
     out1 = TorchBatchNorm2D(axis=1,  momentum=0.1, epsilon=1e-05, weights=t2kw_bn2d(srctorch.bn1) )(input)
     out1 = KL.Activation( keras.backend.relu ) (out1)
@@ -156,6 +157,11 @@ else:
     model.save (model_path)
     model.save_weights ( os.path.join( os.path.dirname(__file__) , 'weights.h5') )
     
+    model_short = keras.models.Model (_input, outputs[0])
+    model_short.compile ( loss='mse', optimizer='adam' )
+    model_short.save ( os.path.join( os.path.dirname(__file__) , "2DFAN-4_light.h5" ) )
+    model_short.save_weights ( os.path.join( os.path.dirname(__file__) , '_light_weights.h5') )
+    
 def transform(point, center, scale, resolution, invert=False):
     _pt = torch.ones(3)
     _pt[0] = point[0]
@@ -268,7 +274,7 @@ pts_img = pts_img.view(68, 2).numpy()
 m_pts_img = get_preds_fromhm2 (m_out[0], center, scale)
 
 print ('pts1 == pts2 == %s' % ( np.array_equal(pts_img, m_pts_img) ) )
-
+import code
 code.interact(local=dict(globals(), **locals()))
 
 #print ( np.array_equal (fa_out, m_out) ) #>>> False
diff --git a/facelib/2DFAN-4.h5 b/facelib/2DFAN-4.h5
index 8c5079a..9d995f9 100644
Binary files a/facelib/2DFAN-4.h5 and b/facelib/2DFAN-4.h5 differ
diff --git a/facelib/LandmarksExtractor.py b/facelib/LandmarksExtractor.py
index 640b1b8..210fd5c 100644
--- a/facelib/LandmarksExtractor.py
+++ b/facelib/LandmarksExtractor.py
@@ -121,7 +121,7 @@ class LandmarksExtractor(object):
             
             predicted = self.keras_model.predict (image)
                 
-            pts_img = get_pts_from_predict ( predicted[-1][0], center, scale)
+            pts_img = get_pts_from_predict ( predicted[-1], center, scale)
             pts_img = [ ( int(pt[0]), int(pt[1]) ) for pt in pts_img ]             
             landmarks.append ( ( (left, top, right, bottom),pts_img ) )