feat sketch 提取接口

fix
2024-08-14 17:43:19 +08:00
parent af3c7c2e51
commit 8b876b103f
1 changed files with 3 additions and 0 deletions
--- a/app/service/image2sketch/server.py
+++ b/app/service/image2sketch/server.py
@@ -55,6 +55,9 @@ class Image2SketchServer:
        self.data['A'] = transform(A)
        self.data['A'] = self.data['A'].unsqueeze(0).to(device)

+    def __del__(self):
+        torch.cuda.empty_cache()
+
    def get_result(self):
        self.model.set_input(self.data)
        self.model.test()  # run inference