Update.

diff --git a/main.py b/main.py
index 21666d1..5dceefc 100755 (executable)
--- a/main.py
+++ b/main.py
@@ -399,7 +399,7 @@ def predict_the_four_grids(
     if with_noise:
         imt_set = add_noise_imt(imt_set)
 
     if with_noise:
         imt_set = add_noise_imt(imt_set)
 

-    result = ae_predict(model, imt_set, local_device=local_device, desc=None)
+    result = ae_predict(model, imt_set, local_device=local_device)

     result = (result * masks).reshape(-1, 4, result.size(1)).sum(dim=1)
 
     return result
     result = (result * masks).reshape(-1, 4, result.size(1)).sum(dim=1)
 
     return result