Merge pull request #10 from jin530/patch-1

Update wide_deep.py

pytorch_widedeep/models/wide_deep.py

@@ -131,7 +131,11 @@ class WideDeep(nn.Module):
 
         if self.deephead is None:
             if head_layers is not None:
-                input_dim: int = self.deepdense.output_dim + self.deeptext.output_dim + self.deepimage.output_dim  # type:ignore
+                input_dim: int = self.deepdense.output_dim # type:ignore
+                if self.deeptext is not None:
+                    input_dim += self.deeptext.output_dim 
+                if self.deepimage is not None:
+                    input_dim += self.deepimage.output_dim  
                 head_layers = [input_dim] + head_layers
                 if not head_dropout:
                     head_dropout = [0.0] * (len(head_layers) - 1)