X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=attentiontoy1d.py;h=6540a0f03bd36316bcd875a90058b3b831dff545;hb=b27b7cc54f450bb5fe8c9ea2faf5e01d0082889a;hp=ad0c0b142a15007a1af64f1b4148056ecf78cf59;hpb=602df813edceaae1b4eb2c69bb5cf0f5823444a7;p=pytorch.git

diff --git a/attentiontoy1d.py b/attentiontoy1d.py
index ad0c0b1..6540a0f 100755
--- a/attentiontoy1d.py
+++ b/attentiontoy1d.py
@@ -10,9 +10,11 @@ import torch, math, sys, argparse
 from torch import nn
 from torch.nn import functional as F
 
+import matplotlib.pyplot as plt
+
 ######################################################################
 
-parser = argparse.ArgumentParser(description='Toy RNN.')
+parser = argparse.ArgumentParser(description='Toy attention model.')
 
 parser.add_argument('--nb_epochs',
                     type = int, default = 250)
@@ -146,9 +148,6 @@ def generate_sequences(nb):
 
 ######################################################################
 
-import matplotlib.pyplot as plt
-import matplotlib.collections as mc
-
 def save_sequence_images(filename, sequences, tr = None, bx = None):
     fig = plt.figure()
     ax = fig.add_subplot(1, 1, 1)