Modified torch module classes to be internel

bonsai-rx · Jan 29, 2025 · 139488e · 139488e
1 parent d06a4c2
commit 139488e
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/src/Bonsai.ML.Torch/NeuralNets/Models/AlexNet.cs b/src/Bonsai.ML.Torch/NeuralNets/Models/AlexNet.cs
@@ -7,7 +7,7 @@ namespace Bonsai.ML.Torch.NeuralNets.Models
     /// <summary>
     /// Modified version of original AlexNet to fix CIFAR10 32x32 images.
     /// </summary>
-    public class AlexNet : Module<Tensor, Tensor>
+    internal class AlexNet : Module<Tensor, Tensor>
     {
         private readonly Module<Tensor, Tensor> features;
         private readonly Module<Tensor, Tensor> avgPool;

diff --git a/src/Bonsai.ML.Torch/NeuralNets/Models/MNIST.cs b/src/Bonsai.ML.Torch/NeuralNets/Models/MNIST.cs
@@ -7,7 +7,7 @@ namespace Bonsai.ML.Torch.NeuralNets.Models
     /// <summary>
     /// Represents a simple convolutional neural network for the MNIST dataset.
     /// </summary>
-    public class MNIST : Module<Tensor,Tensor>
+    internal class MNIST : Module<Tensor,Tensor>
     {
         private readonly Module<Tensor, Tensor> conv1;
         private readonly Module<Tensor, Tensor> conv2;

diff --git a/src/Bonsai.ML.Torch/NeuralNets/Models/MobileNet.cs b/src/Bonsai.ML.Torch/NeuralNets/Models/MobileNet.cs
@@ -9,7 +9,7 @@ namespace Bonsai.ML.Torch.NeuralNets.Models
     /// <summary>
     /// MobileNet model.
     /// </summary>
-    public class MobileNet : Module<Tensor, Tensor>
+    internal class MobileNet : Module<Tensor, Tensor>
     {
         private readonly long[] planes = [ 64, 128, 128, 256, 256, 512, 512, 512, 512, 512, 512, 1024, 1024 ];
         private readonly long[] strides = [ 1, 2, 1, 2, 1, 2, 1, 1, 1, 1, 1, 2, 1 ];