added squeeze op, optionally reduces batch dim if bs == 1

Author: isaaccorley

torch_enhance/models/base.py

@@ -79,4 +79,5 @@ def enhance(self, x: torch.Tensor) -> torch.Tensor:
         x *= 255.0
         x = x.clamp(0, 255)
         x = x.to(torch.uint8)
+        x = x.squeeze(0)
         return x