diff --git a/models/Conv_DCFE.py b/models/Conv_DCFE.py
index f0a9ba8..8fee4a0 100644
--- a/models/Conv_DCFE.py
+++ b/models/Conv_DCFE.py
@@ -94,12 +94,15 @@ def extra_repr(self):
             ', bases_grad={bases_grad}, mode={mode}, in_channels={in_channels}, out_channels={out_channels}'.format(**self.__dict__)
 
 
-
-if __name__ == '__main__':
-    layer = Conv_DCFDE(1, 3, kernel_size=3, padding=1, stride=2).cuda()
-    # layer = nn.Conv2d(3, 10, kernel_size=3, padding=1, stride=2).cuda()
-    data = torch.randn(1 , 1, 4, 4).cuda()
-    print(layer(data))
-    print(layer(data))
+###
+### The below lines are commented out because unless Conv_DCFDE doesn't get initiated with coef instance variable, it will return
+### an error at line #68 because they aren't defined. coef variables are D in the paper: See fig 1. Page no - 2.
+###
+#if __name__ == '__main__':
+#    layer = Conv_DCFDE(1, 3, kernel_size=3, padding=1, stride=2).cuda()
+#    # layer = nn.Conv2d(3, 10, kernel_size=3, padding=1, stride=2).cuda()
+#    data = torch.randn(1 , 1, 4, 4).cuda()
+#    print(layer(data))
+#    print(layer(data))
 
 
diff --git a/models/resnet18_dcf_bsensemble_imgnet.py b/models/resnet18_dcf_bsensemble_imgnet.py
index 16a40f1..ea9f078 100644
--- a/models/resnet18_dcf_bsensemble_imgnet.py
+++ b/models/resnet18_dcf_bsensemble_imgnet.py
@@ -112,7 +112,7 @@ def __init__(self, block, num_blocks, num_classes=10, num_bases=12, num_member=1
             self.coeff_list = coeff_list
         
 
-    def add_branch(self, num_outputs, block_expand=0):
+    def add_branch(self, num_outputs):
         """
           add a set of new modules that form another branch
         """
@@ -200,4 +200,4 @@ def forward(self, x, task_id=None):
 
 
 def Net(**args):
-    return ResNet(BasicBlock, [2,2,2,2], **args)
\ No newline at end of file
+    return ResNet(BasicBlock, [2,2,2,2], **args)
diff --git a/models/resnet32_dcf_bsensemble.py b/models/resnet32_dcf_bsensemble.py
index 52fc7d7..a8dfd97 100644
--- a/models/resnet32_dcf_bsensemble.py
+++ b/models/resnet32_dcf_bsensemble.py
@@ -103,7 +103,7 @@ def __init__(self, block, num_blocks, num_classes=10, num_bases=12, num_member=1
         for pm in self.coeff_list:
             nn.init.kaiming_normal_(pm)
 
-    def add_branch(self, num_outputs, block_expand=0):
+    def add_branch(self, num_outputs):
         """
           add a set of new modules that form another branch
         """
@@ -187,4 +187,4 @@ def forward(self, x, task_id=None):
 
 
 def Net(**args):
-    return ResNet(BasicBlock, [5,5,5], **args)
\ No newline at end of file
+    return ResNet(BasicBlock, [5,5,5], **args)
diff --git a/train_dcfens_cifar100_CI.py b/train_dcfens_cifar100_CI.py
index 4b4dba7..c532965 100644
--- a/train_dcfens_cifar100_CI.py
+++ b/train_dcfens_cifar100_CI.py
@@ -102,7 +102,7 @@ def train(train_loader, epoch, task, model, total_epoch):
     train_loss = 0
     correct = 0
     total = 0
-    # targets_all = []
+    targets_all = []
 
     previous_cls = sum(class_increments[:task])
     for batch_idx, (inputs, targets) in enumerate(train_loader):
@@ -114,7 +114,7 @@ def train(train_loader, epoch, task, model, total_epoch):
         bs = inputs.shape[0]
         outputs, feat_current = model(torch.cat([inputs] * args.num_member, dim=0), task_id=task)
         outputs = outputs.split(bs)
-        # targets_all.append(targets)
+        targets_all.append(targets)
         loss = torch.sum(torch.stack([criterion(outputs_, targets) for outputs_ in outputs], dim=0), dim=0)
         
         loss.backward()
@@ -412,7 +412,7 @@ def inferecne(test_loader, task, total_task, model):
     ## dataloaders
     train_loader, test_loader = task_data[task][0],task_data[task][1]
     ## add a new network branch
-    net.add_branch(class_increments[task], task)
+    net.add_branch(class_increments[task])
     net.cuda()
 
     #  init model with previous task's params
diff --git a/train_dcfens_imgnet100_CI.py b/train_dcfens_imgnet100_CI.py
index a659290..f313be9 100644
--- a/train_dcfens_imgnet100_CI.py
+++ b/train_dcfens_imgnet100_CI.py
@@ -404,7 +404,7 @@ def inferecne(test_loader, task, total_task, model):
     ## dataloaders
     train_loader, test_loader = task_data[task][0],task_data[task][1]
     ## add a new network branch
-    net.add_branch(class_increments[task], task)
+    net.add_branch(class_increments[task])
     net.cuda()
 
     #  init model with previous task's params
@@ -466,4 +466,4 @@ def inferecne(test_loader, task, total_task, model):
 
 print()
 print(ci_acc_list)
-print('\n For Class-Incremental Learning, Average Incremental Accuracy: {:.2f}'.format(np.mean(ci_acc_list)))
\ No newline at end of file
+print('\n For Class-Incremental Learning, Average Incremental Accuracy: {:.2f}'.format(np.mean(ci_acc_list)))