Move monitors to device

BindsNET · Feb 17, 2021 · 09b16ae · 09b16ae
1 parent 5ee9497
commit 09b16ae
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/examples/mnist/supervised_mnist.py b/examples/mnist/supervised_mnist.py
@@ -107,8 +107,8 @@
     network.to("cuda")
 
 # Voltage recording for excitatory and inhibitory layers.
-exc_voltage_monitor = Monitor(network.layers["Ae"], ["v"], time=time)
-inh_voltage_monitor = Monitor(network.layers["Ai"], ["v"], time=time)
+exc_voltage_monitor = Monitor(network.layers["Ae"], ["v"], time=time, device=device)
+inh_voltage_monitor = Monitor(network.layers["Ai"], ["v"], time=time, device=device)
 network.add_monitor(exc_voltage_monitor, name="exc_voltage")
 network.add_monitor(inh_voltage_monitor, name="inh_voltage")