diff --git a/convert_onnx.py b/convert_onnx.py
new file mode 100644
index 0000000..b1745ac
--- /dev/null
+++ b/convert_onnx.py
@@ -0,0 +1,41 @@
+import os
+import torch
+from modules.xfeat import XFeat
+from modules.xfeat import XFeatModel
+import onnxruntime as ort
+
+# os.environ['CUDA_VISIBLE_DEVICES'] = '' #Force CPU, comment for GPU
+
+# set the model to evaluation mode
+net = XFeatModel().eval()
+
+# load the pretrained weights
+net.load_state_dict(torch.load("weights/xfeat.pt", map_location=torch.device('cpu')))
+
+# Random input
+x = torch.randn(1, 1, 640, 640)
+
+# export to ONNX
+torch.onnx.export(net, x, "xfeat.onnx", verbose=True,
+                  input_names=['input'],
+                  output_names=['output_feats', "output_keypoints", "output_heatmap"],
+                  opset_version=11)
+
+print("ONNX model saved as xfeat.onnx")
+
+# check the onnx model with onnxruntime
+ort_session = ort.InferenceSession("xfeat.onnx")
+print("ONNX model loaded successfully")
+
+outputs = ort_session.run(None, {"input": x.numpy()})
+
+# pytorch model outputs
+torch_outputs = net(x)
+
+# compare the outputs
+for i in range(len(outputs)):
+    print(f"onnx output shape {i}: {outputs[i].shape}")
+    print(f"torch output shape {i}: {torch_outputs[i].shape}")
+    print(f"Output {i} comparison: {torch.allclose(torch_outputs[i], torch.tensor(outputs[i]))}")
+    print(f'Output {i} max diff: {torch.max(torch.abs(torch_outputs[i] - torch.tensor(outputs[i])))}')
+    print("\n")
diff --git a/detect_example.py b/detect_example.py
new file mode 100644
index 0000000..28feb7e
--- /dev/null
+++ b/detect_example.py
@@ -0,0 +1,40 @@
+import os
+import numpy as np
+import torch
+import tqdm
+import cv2
+
+from modules.xfeat import XFeat
+
+# os.environ['CUDA_VISIBLE_DEVICES'] = '' #Force CPU, comment for GPU
+
+xfeat = XFeat()
+
+# read image
+img = cv2.imread('D:/work/xfeatc/data/1.png', cv2.IMREAD_GRAYSCALE)
+
+# convert to tensor [1, 1, H, W]
+x = torch.tensor(img, dtype = torch.float32)
+x = x[None, None, :, :] / 255.0
+
+print(x.shape)
+
+points = xfeat.detectAndCompute(x, top_k = 2)[0]
+
+print("----------------")
+
+# draw on image
+cimg = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
+for p in points['keypoints']:
+    cv2.circle(cimg, (int(p[0]), int(p[1])), 3, (255, 0, 0), -1)
+
+cv2.imshow('image', cimg)
+cv2.waitKey(0)
+
+# print keypoints
+for p in points['keypoints']:
+    print("{}, {}".format(p[0], p[1]))
+
+# print descriptors
+for d in points['descriptors']:
+    print(d)
\ No newline at end of file
diff --git a/modules/interpolator.py b/modules/interpolator.py
index 699c1b8..d7594aa 100644
--- a/modules/interpolator.py
+++ b/modules/interpolator.py
@@ -29,5 +29,5 @@ def forward(self, x, pos, H, W):
             [B, N, C] sampled channels at 2d positions
         """
         grid = self.normgrid(pos, H, W).unsqueeze(-2).to(x.dtype)
-        x = F.grid_sample(x, grid, mode = self.mode , align_corners = False)
+        x = F.grid_sample(x, grid, mode = self.mode , align_corners = self.align_corners)
         return x.permute(0,2,3,1).squeeze(-2)
\ No newline at end of file
diff --git a/modules/model.py b/modules/model.py
index 57539fd..4ddb8e2 100644
--- a/modules/model.py
+++ b/modules/model.py
@@ -125,30 +125,34 @@ def forward(self, x):
 			input:
 				x -> torch.Tensor(B, C, H, W) grayscale or rgb images
 			return:
-				feats     ->  torch.Tensor(B, 64, H/8, W/8) dense local features
-				keypoints ->  torch.Tensor(B, 65, H/8, W/8) keypoint logit map
+				feats     ->  torch.Tensor(B, H/8, W/8, 64) dense local features
+				keypoints ->  torch.Tensor(B, H/8, W/8, 65) keypoint logit map
 				heatmap   ->  torch.Tensor(B,  1, H/8, W/8) reliability map
 
 		"""
 		#dont backprop through normalization
 		with torch.no_grad():
 			x = x.mean(dim=1, keepdim = True)
-			x = self.norm(x)
+			# x = self.norm(x)
 
 		#main backbone
-		x1 = self.block1(x)
-		x2 = self.block2(x1 + self.skip1(x))
-		x3 = self.block3(x2)
-		x4 = self.block4(x3)
-		x5 = self.block5(x4)
+		x1 = self.block1(x)								# x1: [B, 24, H/4, W/4]
+		x2 = self.block2(x1 + self.skip1(x))				# x2: [B, 24, H/4, W/4]
+		x3 = self.block3(x2)							# x3: [B, 64, H/8, W/8]
+		x4 = self.block4(x3)							# x4: [B, 64, H/16, W/16]
+		x5 = self.block5(x4)							# x5: [B, 64, H/32, W/32]
 
 		#pyramid fusion
-		x4 = F.interpolate(x4, (x3.shape[-2], x3.shape[-1]), mode='bilinear')
-		x5 = F.interpolate(x5, (x3.shape[-2], x3.shape[-1]), mode='bilinear')
-		feats = self.block_fusion( x3 + x4 + x5 )
+		x4 = F.interpolate(x4, (x3.shape[-2], x3.shape[-1]), mode='bilinear')  # x4: [B, 64, H/8, W/8]
+		x5 = F.interpolate(x5, (x3.shape[-2], x3.shape[-1]), mode='bilinear')  # x5: [B, 64, H/8, W/8]
+		feats = self.block_fusion( x3 + x4 + x5 )  # feats: [B, 64, H/8, W/8]
 
 		#heads
 		heatmap = self.heatmap_head(feats) # Reliability map
 		keypoints = self.keypoint_head(self._unfold2d(x, ws=8)) #Keypoint map logits
 
-		return feats, keypoints, heatmap
+		# convert feats and keypoints from [B, C, H, W] to [B, H, W, C]
+		out_feats = feats.permute(0, 2, 3, 1)
+		out_keypoints = keypoints.permute(0, 2, 3, 1)
+
+		return out_feats, out_keypoints, heatmap