diff --git a/bob/ip/binseg/models/driu_od.py b/bob/ip/binseg/models/driu_od.py
index e45755afb6bcff8f708b99220d683b84148427be..5ba6a7e27e38b5a89a23d72a37fae9d374f516ac 100644
--- a/bob/ip/binseg/models/driu_od.py
+++ b/bob/ip/binseg/models/driu_od.py
@@ -24,7 +24,9 @@ class DRIUOD(torch.nn.Module):
 
     def __init__(self, in_channels_list=None):
         super(DRIUOD, self).__init__()
-        in_upsample2, in_upsample_4, in_upsample_8, in_upsample_16 = in_channels_list
+        in_upsample2, in_upsample_4, in_upsample_8, in_upsample_16 = (
+            in_channels_list
+        )
 
         self.upsample2 = UpsampleCropBlock(in_upsample2, 16, 4, 2, 0)
         # Upsample layers
@@ -83,14 +85,16 @@ def driu_od(pretrained_backbone=True, progress=True):
     """
 
     backbone = vgg16_for_segmentation(
-        pretrained=pretrained_backbone, progress=progress,
-        return_features=[3, 8, 14, 22],
+        pretrained=pretrained_backbone,
+        progress=progress,
+        return_features=[8, 14, 22, 29],
     )
     head = DRIUOD([128, 256, 512, 512])
 
     order = [("backbone", backbone), ("head", head)]
     if pretrained_backbone:
         from .normalizer import TorchVisionNormalizer
+
         order = [("normalizer", TorchVisionNormalizer())] + order
 
     model = torch.nn.Sequential(OrderedDict(order))