Update.

diff --git a/main.py b/main.py
index d8dffe2..c6eedfb 100755 (executable)
--- a/main.py
+++ b/main.py
@@ -380,8 +380,12 @@ def add_noise(imt_set):
     return torch.cat([input[:, None], masks[:, None], targets[:, None]], dim=1)
 
 
+######################################################################
+
 # IMT for input / masks / target
 
+# Generate a batch for prediction
+
 
 def batch_for_prediction_imt(input):
     nb = input.size(0)