From b50d68de02beafbf8bc9e5e5dd2b39da8c80e997 Mon Sep 17 00:00:00 2001
From: Nicolas Kaenzig <nkaenzig@gmail.com>
Date: Wed, 30 Oct 2024 13:24:04 +0100
Subject: [PATCH] convert tensor to list

---
 .../vision/models/networks/decoders/segmentation/decoder2d.py   | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/eva/vision/models/networks/decoders/segmentation/decoder2d.py b/src/eva/vision/models/networks/decoders/segmentation/decoder2d.py
index 36c96410..c43b351c 100644
--- a/src/eva/vision/models/networks/decoders/segmentation/decoder2d.py
+++ b/src/eva/vision/models/networks/decoders/segmentation/decoder2d.py
@@ -50,6 +50,8 @@ def _forward_features(self, features: torch.Tensor | List[torch.Tensor]) -> torc
             A tensor of shape (batch_size, hidden_size, n_patches_height,
             n_patches_width) which is feature map of the decoder head.
         """
+        if isinstance(features, torch.Tensor):
+            features = [features]
         if not isinstance(features, list) or features[0].ndim != 4:
             raise ValueError(
                 "Input features should be a list of four (4) dimensional inputs of "