Refactor channel determination logic in TorchSharp.

Author: AhmedZero

src/TorchVision/Functional.cs

@@ -37,6 +37,17 @@ private static bool AssertTensorImage(Tensor img)
                     return true;
                 }
 
+                public static long GetImageNumChannels(Tensor img)
+                {
+                    AssertTensorImage(img);
+                    var ndim_ = img.ndim;
+                    return ndim_ switch {
+                        2 => 1,
+                        > 2 => img.shape[ndim_ - 3],
+                        _ => throw new ArgumentException($"Input ndim should be 2 or more. Got {ndim_}"),
+                    };
+                }
+
                 /// <summary>
                 /// Get the image dimensions
                 /// </summary>

src/TorchVision/Normalize.cs

@@ -23,7 +23,7 @@ internal Normalize(double[] means, double[] stdevs,bool inplace = false)
 
             public Tensor call(Tensor input)
             {
-                var expectedChannels = (input.ndim == 4) ? input.size(1) : input.size(0);
+                var expectedChannels = transforms.functional.GetImageNumChannels(input);
                 if (expectedChannels != means.Length)
                     throw new ArgumentException("The number of channels is not equal to the number of means and standard deviations");
                 return transforms.functional.normalize(input, means, stdevs, inplace);