diff --git a/IGEV-Stereo/core/geometry.py b/IGEV-Stereo/core/geometry.py
index 12f1a85..dda4aa1 100644
--- a/IGEV-Stereo/core/geometry.py
+++ b/IGEV-Stereo/core/geometry.py
@@ -60,6 +60,7 @@ class Combined_Geo_Encoding_Volume:
     
     @staticmethod
     def corr(fmap1, fmap2):
+        # batch, dim, ht, wd
         B, D, H, W1 = fmap1.shape
         _, _, _, W2 = fmap2.shape
         fmap1 = fmap1.view(B, D, H, W1)
diff --git a/IGEV-Stereo/core/igev_stereo.py b/IGEV-Stereo/core/igev_stereo.py
index 18d31ef..cb02e68 100644
--- a/IGEV-Stereo/core/igev_stereo.py
+++ b/IGEV-Stereo/core/igev_stereo.py
@@ -167,6 +167,7 @@ class IGEVStereo(nn.Module):
             match_right = self.desc(self.conv(features_right[0]))
             gwc_volume = build_gwc_volume(match_left, match_right, self.args.max_disp//4, 8)
             gwc_volume = self.corr_stem(gwc_volume)
+            # 3d unet
             gwc_volume = self.corr_feature_att(gwc_volume, features_left[0])
             geo_encoding_volume = self.cost_agg(gwc_volume, features_left)