disposable pattern

AutoencoderKL.cs

@@ -1,8 +1,9 @@
+using System;
 using TorchSharp;
 
-public class AutoencoderKL
+public class AutoencoderKL : IDisposable
 {
-    private readonly torch.jit.ScriptModule _model;
+    private torch.jit.ScriptModule _model;
     private readonly float _scale;
     public torch.Device Device {get;}
 
@@ -21,4 +22,10 @@ public class AutoencoderKL
         tokenTensor = 1.0f / _scale * tokenTensor;
         return (torch.Tensor)_model.forward(tokenTensor);
     }
+
+    public void Dispose()
+    {
+        _model.Dispose();
+        _model = null;
+    }
 }

ClipEnocder.cs

@@ -1,8 +1,9 @@
+using System;
 using TorchSharp;
 
-public class ClipEncoder
+public class ClipEncoder : IDisposable
 {
-    private readonly torch.jit.ScriptModule _model;
+    private torch.jit.ScriptModule _model;
     public torch.Device Device {get;}
 
     public ClipEncoder(string modelPath, torch.Device device)
@@ -17,4 +18,10 @@ public class ClipEncoder
     {
         return (torch.Tensor)_model.forward(tokenTensor);
     }
+
+    public void Dispose()
+    {
+        _model.Dispose();
+        _model = null;
+    }
 }

DDPM.cs

@@ -1,7 +1,7 @@
 using System;
 using TorchSharp;
 
-public class DDPM
+public class DDPM : IDisposable
 {
     private readonly torch.jit.ScriptModule _model;
     public torch.Device Device {get;}
@@ -36,4 +36,10 @@ public class DDPM
     {
         return _model.invoke<torch.Tensor>("predict_start_from_z_and_v", z, t, v);
     }
+
+    public void Dispose()
+    {
+        _model.Dispose();
+        _model = null;
+    }
 }

Program.cs

@@ -6,9 +6,6 @@ using TorchSharp;
 
 torchvision.io.DefaultImager = new torchvision.io.SkiaImager();
 var device = TorchSharp.torch.device("cuda:0");
-var ddpm = new DDPM("ddim_v_sampler.ckpt", device);
-var ddimSampler = new DDIMSampler(ddpm);
-var autoencoderKL = new AutoencoderKL("autoencoder_kl.ckpt", device);
 var clipEncoder = new ClipEncoder("clip_encoder.ckpt", device);
 var start_token = 49406;
 var end_token = 49407;
@@ -21,7 +18,7 @@ var dictionary = new Dictionary<string, long>(){
     {"green", 1901},
 };
 
-var batch = 2;
+var batch = 1;
 
 var prompt = "a wild cute green cat";
 var tokens = prompt.Split(' ').Select(x => dictionary[x]).ToList();
@@ -33,15 +30,23 @@ var tokenTensor = torch.tensor(tokens.ToArray(), dtype: torch.ScalarType.Int64,
 tokenTensor = tokenTensor.repeat(batch, 1);
 var unconditional_tokenTensor = torch.tensor(uncontional_tokens.ToArray(), dtype: torch.ScalarType.Int64, device: device);
 unconditional_tokenTensor = unconditional_tokenTensor.repeat(batch, 1);
-var img = torch.randn(batch, 4, 96, 96, dtype: torch.ScalarType.Float32, device: device);
+var img = torch.randn(batch, 4, 64, 64, dtype: torch.ScalarType.Float32, device: device);
 var t = torch.full(new[]{batch, 1L}, value: batch, dtype: torch.ScalarType.Int32, device: device);
 var condition = clipEncoder.Forward(tokenTensor);
 var unconditional_condition = clipEncoder.Forward(unconditional_tokenTensor);
+
+clipEncoder.Dispose();
+var ddpm = new DDPM("ddim_v_sampler.ckpt", device);
+var ddimSampler = new DDIMSampler(ddpm);
 var ddim_steps = 50;
 img = ddimSampler.Sample(img, condition, unconditional_condition, ddim_steps);
+ddpm.Dispose();
+
+var autoencoderKL = new AutoencoderKL("autoencoder_kl.ckpt", device);
 var decoded_images = (torch.Tensor)autoencoderKL.Forward(img);
 decoded_images = torch.clamp((decoded_images + 1.0) / 2.0, 0.0, 1.0);
 
+
 for(int i = 0; i!= batch; ++i)
 {
     var image = decoded_images[i];