From aaf6fa3ea6c8552ac9de0be8143907150d04dd7e Mon Sep 17 00:00:00 2001
From: andreABbauer <62743847+andreABbauer@users.noreply.github.com>
Date: Thu, 7 May 2020 15:57:03 +0200
Subject: [PATCH] Update comment to match code

I am checking this out with the paper and a few experiments on the side and i think the comment was wrong. Can someone please confirm?
---
 lib/models/seg_hrnet_ocr.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lib/models/seg_hrnet_ocr.py b/lib/models/seg_hrnet_ocr.py
index ed9df629..8ab88b44 100644
--- a/lib/models/seg_hrnet_ocr.py
+++ b/lib/models/seg_hrnet_ocr.py
@@ -62,7 +62,7 @@ def forward(self, feats, probs):
         feats = feats.permute(0, 2, 1) # batch x hw x c 
         probs = F.softmax(self.scale * probs, dim=2)# batch x k x hw
         ocr_context = torch.matmul(probs, feats)\
-        .permute(0, 2, 1).unsqueeze(3)# batch x k x c
+        .permute(0, 2, 1).unsqueeze(3) # batch x c x k x 1
         return ocr_context