diff --git a/model.py b/model.py
index f15a9fe..73d9d73 100644
--- a/model.py
+++ b/model.py
@@ -71,18 +71,16 @@ class FastRCNN(nn.Module):
         # hidden_dim -> hidden_dim.                                                  #
         ##############################################################################
         # Replace "pass" statement with your code
-        self.cls_head = nn.Sequential(
+        self.shared_fc = nn.Sequential(
             nn.Linear(in_dim, hidden_dim),
             nn.Dropout(drop_ratio),
             nn.ReLU(),
-            nn.Linear(hidden_dim, num_classes+1)
-        )
-        self.bbox_head = nn.Sequential(
-            nn.Linear(in_dim, hidden_dim),
-            nn.Dropout(drop_ratio),
-            nn.ReLU(),
-            nn.Linear(hidden_dim, 4)
+            nn.Linear(hidden_dim, hidden_dim)
         )
+
+        self.cls_head = nn.Linear(hidden_dim, self.num_classes+1) # The cls head is a Linear layer that predicts num_classes + 1 (background).
+
+        self.bbox_head = nn.Linear(hidden_dim, 4)# The det head is a Linear layer that predicts offsets(dim=4).
         ##############################################################################
         #                               END OF YOUR CODE                             #
         ##############################################################################
@@ -139,8 +137,9 @@ class FastRCNN(nn.Module):
         # print(feat.shape)
 
         # forward heads, get predicted cls scores & offsets
-        cls_scores=self.cls_head(feat)
-        bbox_offsets=self.bbox_head(feat)
+        shared_feat = self.shared_fc(feat)
+        cls_scores=self.cls_head(shared_feat)
+        bbox_offsets=self.bbox_head(shared_feat)
         # print(cls_scores.shape, bbox_offsets.shape)
 
         # assign targets with proposals
@@ -216,11 +215,11 @@ class FastRCNN(nn.Module):
         # perform RoI Pool & mean pool
         feat=torchvision.ops.roi_pool(feat, torch.cat((proposal_batch_ids.unsqueeze(1), proposals),dim=1), output_size=(self.roi_output_w, self.roi_output_h))
         feat = feat.mean(dim=[2, 3])
-
+        shared_feat = self.shared_fc(feat)
         # forward heads, get predicted cls scores & offsets
-        cls_scores = self.cls_head(feat)
+        cls_scores = self.cls_head(shared_feat)
         # print(cls_scores.shape)
-        bbox_offsets = self.bbox_head(feat)
+        bbox_offsets = self.bbox_head(shared_feat)
         # print(bbox_offsets.shape)
         # get predicted boxes & class label & confidence probability
         proposals = generate_proposal(proposals, bbox_offsets)