Update.

[pytorch.git] / pscan.py

diff --git a/pscan.py b/pscan.py
index 36490ff..d009692 100755 (executable)
--- a/pscan.py
+++ b/pscan.py
@@ -1,11 +1,6 @@
 #!/usr/bin/env python
 
-import math
-
-import torch, torchvision
-
-from torch import nn
-from torch.nn import functional as F
+import torch
 
 ######################################################################
 
@@ -23,7 +18,9 @@ def naive_rec(A, X, Y0):
 
 ######################################################################
 
-# A is NxTx1 and X is NxTxD
+# A is NxTx1
+# X is NxTxD
+# Y0 is NxD
 #
 # Returns Y defined with
 #