1 files changed, 7 insertions, 4 deletions
diff --git a/code/weightvis.js b/code/weightvis.js
index 6e9fc1d..428d3ce 100644
--- a/code/weightvis.js
+++ b/code/weightvis.js
@@ -510,8 +510,8 @@ var NetVisualizer = (function() {
         this.reposition();
     };
 
-    NetVisualizer.prototype.visualActivate = function(inputs) {
-        var activations = this.activate(inputs);
+    NetVisualizer.prototype.visualActivate = function(inputs, step) {
+        var activations = this.activate(inputs, step);
         var flatActivations = [].concat.apply([], activations);
         var max = d3.max(flatActivations, Math.abs),
             min = d3.min(flatActivations, Math.abs);
@@ -521,10 +521,13 @@ var NetVisualizer = (function() {
             .style("fill", fNodeColor);
     };
 
-    NetVisualizer.prototype.activate = function(inputs) {
+    NetVisualizer.prototype.activate = function(inputs, step) {
         // assumes the sigmoid function is used as activation
         function activation(x) { return 1 / (1 + Math.exp(-x)); }
-        var weights = this.data[this.layout.data.batches-1][this.layout.data.epochs-1];
+        if (typeof(step) == "undefined") step = this.getSteps() - 1;
+        var batch = Math.floor(step / this.layout.data.epochs),
+            epoch = step % this.layout.data.epochs;
+        var weights = this.data[batch][epoch];
         var lastLayerOut = inputs;
         var activations = [inputs];
         var hasBias = false;