Removed the following files & folders:
1. assets
2. configs
3. data
4. notebooks
5. scripts
6. environment_train.yaml
7. train.py
8. test.py

Replaced

        feat_c0 = rearrange(feat_c0, 'n c h w -> n (h w) c')
        feat_c1 = rearrange(feat_c1, 'n c h w -> n (h w) c')
to 
        feat_c0 = feat_c0.permute(0, 2, 3, 1).flatten(1, 2)
        feat_c1 = feat_c1.permute(0, 2, 3, 1).flatten(1, 2)
in src/loftr.py: ONNX export compatibility


Removed         
        # detect NaN during mixed precision training
        if self.config['replace_nan'] and (torch.any(torch.isnan(feat_c0)) or torch.any(torch.isnan(feat_c1))):
            detect_NaN(feat_c0, feat_c1)
from src/loftr.py: Removed in inference


Removed
        # 4. Random sampling of training samples for fine-level LoFTR
        # (optional) pad samples with gt coarse-level matches
        if self.training:
            # NOTE:
            # The sampling is performed across all pairs in a batch without manually balancing
            # #samples for fine-level increases w.r.t. batch_size
            if 'mask0' not in data:
                num_candidates_max = mask.size(0) * max(
                    mask.size(1), mask.size(2))
            else:
                num_candidates_max = compute_max_candidates(
                    data['mask0'], data['mask1'])
            num_matches_train = int(num_candidates_max *
                                    self.train_coarse_percent)
            num_matches_pred = len(b_ids)
            assert self.train_pad_num_gt_min < num_matches_train, "min-num-gt-pad should be less than num-train-matches"

            # pred_indices is to select from prediction
            if num_matches_pred <= num_matches_train - self.train_pad_num_gt_min:
                pred_indices = torch.arange(num_matches_pred, device=_device)
            else:
                pred_indices = torch.randint(
                    num_matches_pred,
                    (num_matches_train - self.train_pad_num_gt_min, ),
                    device=_device)

            # gt_pad_indices is to select from gt padding. e.g. max(3787-4800, 200)
            gt_pad_indices = torch.randint(
                    len(data['spv_b_ids']),
                    (max(num_matches_train - num_matches_pred,
                        self.train_pad_num_gt_min), ),
                    device=_device)
            mconf_gt = torch.zeros(len(data['spv_b_ids']), device=_device)  # set conf of gt paddings to all zero

            b_ids, i_ids, j_ids, mconf = map(
                lambda x, y: torch.cat([x[pred_indices], y[gt_pad_indices]],
                                       dim=0),
                *zip([b_ids, data['spv_b_ids']], [i_ids, data['spv_i_ids']],
                     [j_ids, data['spv_j_ids']], [mconf, mconf_gt]))
from src/loftr.py: Training specific code

Removed the code of scale norm from coarse_matching.py
        # prevent fp16 overflow during mixed precision training
        feat_c0, feat_c1 = map(lambda feat: feat / feat.shape[-1]**.5, [feat_c0, feat_c1])
because the loftr was doing it again after that. Now doing single scale norm that is before coarse_matching call in loftr forward.

Removed 
        # detect NaN during mixed precision training
        if self.config['replace_nan'] and (torch.any(torch.isnan(feat_f0_unfold)) or torch.any(torch.isnan(feat_f1_unfold))):
            detect_NaN(feat_f0_unfold, feat_f1_unfold)
from loftr.py: Removed in inference

Removed
        if 'mask0' in data:
            mask_c0, mask_c1 = data['mask0'], data['mask1']
from loftr.py: Not required at inference


Removed else from
        if data['hw0_i'] == data['hw1_i']:  # faster & better BN convergence
            # backbone((2N, 1, H, W)) -> {feats_c: (2N, C, H/8, W/8), feats_x2: (2N, C2, H/4, W/4), feats_x1: (2N, C1, H/2, W/2)}
            ret_dict = self.backbone(torch.cat([data['image0'], data['image1']], dim=0))
            
            feats_c = ret_dict['feats_c']
            
            data.update({
                'feats_x2': ret_dict['feats_x2'],
                'feats_x1': ret_dict['feats_x1'],
            })

            # feat_c0: (N, C0, H0/8, W0/8), feat_c1: (N, C1, H1/8, W1/8)
            (feat_c0, feat_c1) = feats_c.split(data['bs'])
        else:  # handle different input shapes
            # backbone((N, 1, H, W)) -> {feats_c: (N, C, H/8, W/8), feats_x2: (N, C2, H/4, W/4), feats_x1: (N, C1, H/2, W/2)}
            ret_dict0, ret_dict1 = self.backbone(data['image0']), self.backbone(data['image1'])

            feat_c0 = ret_dict0['feats_c']
            feat_c1 = ret_dict1['feats_c']

            data.update({
                'feats_x2_0': ret_dict0['feats_x2'],
                'feats_x1_0': ret_dict0['feats_x1'],
                'feats_x2_1': ret_dict1['feats_x2'],
                'feats_x1_1': ret_dict1['feats_x1'],
            })
from loftr.py: Condition if branch based on input properties will break onnx tracing or create unnecessarily complex graph. Std. production pipelines resizes both image to exact same dimenions

Removed else from
        if 'mask0' not in data:
            mask_border(mask, self.border_rm, False)
        else:
            mask_border_with_padding(mask, self.border_rm, False, data['mask0'], data['mask1'])
in coarse_matching.py

Removed 
            if mask_c0 is not None:
                sim_matrix = sim_matrix.masked_fill(
                    ~(mask_c0[..., None] * mask_c1[:, None]).bool(),
                    -1e4
                    )
            if mask_c0 is not None:
                    sim_matrix = sim_matrix.float().masked_fill(
                        ~(mask_c0[..., None] * mask_c1[:, None]).bool(),
                        -INF
                        )
from coarse_matching.py as mask is None at inference.

Replaced 
        mask = rearrange(mask, 'b (h0c w0c) (h1c w1c) -> b h0c w0c h1c w1c', **axes_lengths)
        mask_border(mask, self.border_rm, False)
        mask = rearrange(mask, 'b h0c w0c h1c w1c -> b (h0c w0c) (h1c w1c)', **axes_lengths)
from
        mask = mask.view(-1, h0c, w0c, h1c, w1c)
        mask_border(mask, self.border_rm, False)
        mask = mask.view(-1, h0c * w0c, h1c * w1c)
in coarse_matching.py

Replaced
        mkpts0_c = torch.stack([i_ids % data['hw0_c'][1], i_ids // data['hw0_c'][1]], dim=1) * scale0
        mkpts1_c = torch.stack([j_ids % data['hw1_c'][1], j_ids // data['hw1_c'][1]], dim=1) * scale1
with
        w0c_tensor = torch.tensor(w0c, device=conf_matrix.device, dtype=i_ids.dtype)
        w1c_tensor = torch.tensor(w1c, device=conf_matrix.device, dtype=j_ids.dtype)
        
        mkpts0_c = torch.stack([i_ids % w0c_tensor, torch.div(i_ids, w0c_tensor, rounding_mode='trunc')], dim=1) * scale0
        mkpts1_c = torch.stack([j_ids % w1c_tensor, torch.div(j_ids, w1c_tensor, rounding_mode='trunc')], dim=1) * scale1

Removed else block
        else:  # handle different input shapes
            feat_c0, feat_c1 = rearrange(feat_c0, 'b (h w) c -> b c h w', h=data['hw0_c'][0]), rearrange(feat_c1, 'b (h w) c -> b c h w', h=data['hw1_c'][0]) # 1/8 feat
            x2_0, x2_1 = data['feats_x2_0'], data['feats_x2_1'] # 1/4 feat
            x1_0, x1_1 = data['feats_x1_0'], data['feats_x1_1'] # 1/2 feat
            del data['feats_x2_0'], data['feats_x1_0'], data['feats_x2_1'], data['feats_x1_1']

            # 1. fine feature extraction
            feat_f0, feat_f1 = self.inter_fpn(feat_c0, x2_0, x1_0, stride), self.inter_fpn(feat_c1, x2_1, x1_1, stride)

            # 2. unfold(crop) all local windows
            feat_f0 = F.unfold(feat_f0, kernel_size=(W, W), stride=stride, padding=0)
            feat_f0 = rearrange(feat_f0, 'n (c ww) l -> n l ww c', ww=W**2)
            feat_f1 = F.unfold(feat_f1, kernel_size=(W+2, W+2), stride=stride, padding=1)
            feat_f1 = rearrange(feat_f1, 'n (c ww) l -> n l ww c', ww=(W+2)**2)

            # 3. select only the predicted matches
            feat_f0 = feat_f0[data['b_ids'], data['i_ids']]  # [n, ww, cf]
            feat_f1 = feat_f1[data['b_ids'], data['j_ids']]

            return feat_f0, feat_f1
from fine_preprocess.py

Hardcoded to int
            self.W = int(self.config['fine_window_size'])
in fine_preprocess.py

Removed
        if data['b_ids'].shape[0] == 0:
            feat0 = torch.empty(0, self.W**2, self.fine_d_model, device=feat_c0.device)
            feat1 = torch.empty(0, self.W**2, self.fine_d_model, device=feat_c0.device)
            return feat0, feat1
from fine_preprocess.py


Removed stride as params in inter_fpn func in fine_preprocess.py

Replaced 
        feat_f0 = F.unfold(feat_f0, kernel_size=(W, W), stride=stride, padding=0)
            feat_f0 = rearrange(feat_f0, 'n (c ww) l -> n l ww c', ww=W**2)
            feat_f1 = F.unfold(feat_f1, kernel_size=(W+2, W+2), stride=stride, padding=1)
            feat_f1 = rearrange(feat_f1, 'n (c ww) l -> n l ww c', ww=(W+2)**2)
with
        feat_f0 = F.unfold(feat_f0, kernel_size=(W, W), stride=stride, padding=0)
        C, L = feat_f0.shape[1] // (W**2), feat_f0.shape[2]
        feat_f0 = feat_f0.view(-1, C, W**2, L).permute(0, 3, 2, 1)

        feat_f1 = F.unfold(feat_f1, kernel_size=(W+2, W+2), stride=stride, padding=1)
        C_pad, L_pad = feat_f1.shape[1] // ((W+2)**2), feat_f1.shape[2]
        feat_f1 = feat_f1.view(-1, C_pad, (W+2)**2, L_pad).permute(0, 3, 2, 1)
in fine_preprocess.py

Removed
        if M == 0:
            assert self.training == False, "M is always > 0 while training, see coarse_matching.py"
            data.update({
                'conf_matrix_f': torch.empty(0, WW, WW, device=feat_0.device),
                'mkpts0_f': data['mkpts0_c'],
                'mkpts1_f': data['mkpts1_c'],
            })
            return
from fine_matching.py

Removed
        # for fine-level supervision
        if self.training or self.validate:
            data.update({'sim_matrix_ff': conf_matrix_ff})
            data.update({'conf_matrix_f': softmax_matrix_f})
from fine_matching.py

Replaced
        m_ids = torch.arange(M, device=idx_l.device, dtype=torch.long).unsqueeze(-1)
        m_ids = m_ids[:len(data['mconf'])] (Removed)
        idx_r_iids, idx_r_jids = idx_r // W, idx_r % W
with
        m_ids = torch.arange(M, device=idx_l.device, dtype=torch.long)
        idx_r_iids, idx_r_jids = torch.div(idx_r, W, rounding_mode="trunc"), idx_r % W
from fine_matching.py

Removed
        if idx_l.numel() == 0:
            data.update({
                'mkpts0_f': data['mkpts0_c'],
                'mkpts1_f': data['mkpts1_c'],
            })
            return
from fine_matching.py

Replaced
        conf_matrix_ff = conf_matrix_ff.view(M, self.WW, self.W+2, self.W+2)
        conf_matrix_ff = conf_matrix_ff[m_ids, idx_l, idx_r_iids, idx_r_jids]
with
        conf_matrix_ff = conf_matrix_ff.view(-1, self.WW, self.W+2, self.W+2)

        # Flatten multi-dimensional gather for ONNX compatibility
        flat_indices = (m_ids * WW * (W+2) * (W+2)) + (idx_l * (W+2) * (W+2)) + (idx_r_iids * (W+2)) + idx_r_jids
        conf_matrix_ff = conf_matrix_ff.view(-1)
        conf_matrix_ff = conf_matrix_ff[flat_indices.view(-1)]
from fine_matching.py


Replaced
        if data['bs'] == 1:
            scale1 = scale * data['scale1'] if 'scale0' in data else scale
        else:
            scale1 = scale * data['scale1'][data['b_ids']][:len(data['mconf']), ...][:,None,:].expand(-1, -1, 2).reshape(-1, 2) if 'scale0' in data else scale
with
        scale1_tensor = data.get('scale1', torch.tensor(scale, device=feat_0.device))
from fine_matching.py


Replaced
        conf_matrix = conf_matrix.reshape(m, -1)[:len(data['mconf']),...]
from
        conf_matrix = conf_matrix.reshape(m, -1)
from fine_matching.py

