diff --git a/.github/workflows/black.yml b/.github/workflows/black.yml
new file mode 100644
index 00000000..01519d98
--- /dev/null
+++ b/.github/workflows/black.yml
@@ -0,0 +1,11 @@
+name: Black Formater
+
+on: [push, pull_request]
+
+jobs:
+  lint:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v2
+      - uses: actions/setup-python@v2
+      - uses: psf/black@stable
\ No newline at end of file
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index cd425a1c..b450d5fe 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -3,4 +3,4 @@ repos:
     rev: master
     hooks:
     - id: black
-      language_version: python3.8
+      language_version: python3
diff --git a/bindsnet/analysis/plotting.py b/bindsnet/analysis/plotting.py
index 6228d54c..663b93e3 100644
--- a/bindsnet/analysis/plotting.py
+++ b/bindsnet/analysis/plotting.py
@@ -450,7 +450,7 @@ def plot_assignments(
 
             if classes is None:
                 cbar = plt.colorbar(im, cax=cax, ticks=list(range(-1, 11)))
-                cbar.ax.set_yticklabels(["none"] + list(range(10)))
+                cbar.ax.set_yticklabels(["none"] + list(range(11)))
             else:
                 cbar = plt.colorbar(im, cax=cax, ticks=np.arange(-1, len(classes)))
                 cbar.ax.set_yticklabels(["none"] + list(classes))
diff --git a/bindsnet/analysis/visualization.py b/bindsnet/analysis/visualization.py
index 85430717..0b151663 100644
--- a/bindsnet/analysis/visualization.py
+++ b/bindsnet/analysis/visualization.py
@@ -10,7 +10,7 @@ def plot_weights_movie(ws: np.ndarray, sample_every: int = 1) -> None:
     # language=rst
     """
     Create and plot movie of weights.
-    
+
     :param ws: Array of shape ``[n_examples, source, target, time]``.
     :param sample_every: Sub-sample using this parameter.
     """
diff --git a/bindsnet/datasets/alov300.py b/bindsnet/datasets/alov300.py
index c1bb6044..afcb38c1 100644
--- a/bindsnet/datasets/alov300.py
+++ b/bindsnet/datasets/alov300.py
@@ -28,7 +28,7 @@ class ALOV300(Dataset):
     def __init__(self, root, transform, input_size, download=False):
         """
         Class to read the ALOV dataset
-        
+
         :param root: Path to the ALOV folder that contains JPEGImages, Annotations, etc. folders.
         :param input_size: The input size of network that is using this data, for rescaling
         :param download: Specify whether to download the dataset if it is not present
@@ -247,7 +247,7 @@ def _download(self):
         """
         Downloads the correct dataset based on the given parameters
 
-        Relies on self.tag to determine both the name of the folder created for the dataset and for the finding the correct download url. 
+        Relies on self.tag to determine both the name of the folder created for the dataset and for the finding the correct download url.
         """
 
         os.makedirs(self.root)
diff --git a/bindsnet/encoding/encoders.py b/bindsnet/encoding/encoders.py
index 5014875f..e17a4f22 100644
--- a/bindsnet/encoding/encoders.py
+++ b/bindsnet/encoding/encoders.py
@@ -22,7 +22,7 @@ class NullEncoder(Encoder):
     # language=rst
     """
     Pass through of the datum that was input.
-    
+
     .. note::
         This is not a real spike encoder. Be careful with the usage of this class.
     """
diff --git a/bindsnet/evaluation/evaluation.py b/bindsnet/evaluation/evaluation.py
index e653ad4f..d2cf1a08 100644
--- a/bindsnet/evaluation/evaluation.py
+++ b/bindsnet/evaluation/evaluation.py
@@ -30,7 +30,7 @@ def assign_labels(
     n_neurons = spikes.size(2)
 
     if rates is None:
-        rates = torch.zeros(n_neurons, n_labels)
+        rates = torch.zeros((n_neurons, n_labels), device=spikes.device)
 
     # Sum over time dimension (spike ordering doesn't matter).
     spikes = spikes.sum(1)
@@ -112,7 +112,7 @@ def all_activity(
     # Sum over time dimension (spike ordering doesn't matter).
     spikes = spikes.sum(1)
 
-    rates = torch.zeros(n_samples, n_labels)
+    rates = torch.zeros((n_samples, n_labels), device=spikes.device)
     for i in range(n_labels):
         # Count the number of neurons with this label assignment.
         n_assigns = torch.sum(assignments == i).float()
@@ -153,7 +153,7 @@ def proportion_weighting(
     # Sum over time dimension (spike ordering doesn't matter).
     spikes = spikes.sum(1)
 
-    rates = torch.zeros(n_samples, n_labels)
+    rates = torch.zeros((n_samples, n_labels), device=spikes.device)
     for i in range(n_labels):
         # Count the number of neurons with this label assignment.
         n_assigns = torch.sum(assignments == i).float()
@@ -191,7 +191,7 @@ def ngram(
     """
     predictions = []
     for activity in spikes:
-        score = torch.zeros(n_labels)
+        score = torch.zeros(n_labels, device=spikes.device)
 
         # Aggregate all of the firing neurons' indices
         fire_order = []
diff --git a/bindsnet/models/models.py b/bindsnet/models/models.py
index 34158115..a38e94bc 100644
--- a/bindsnet/models/models.py
+++ b/bindsnet/models/models.py
@@ -405,7 +405,9 @@ def __init__(
         w = w / w.max()
         w = (w * self.max_inhib) + self.start_inhib
         recurrent_output_conn = Connection(
-            source=self.layers["Y"], target=self.layers["Y"], w=w,
+            source=self.layers["Y"],
+            target=self.layers["Y"],
+            w=w,
         )
         self.add_connection(recurrent_output_conn, source="Y", target="Y")
 
diff --git a/bindsnet/network/network.py b/bindsnet/network/network.py
index ce61ff89..643b1c40 100644
--- a/bindsnet/network/network.py
+++ b/bindsnet/network/network.py
@@ -197,7 +197,7 @@ def clone(self) -> "Network":
         # language=rst
         """
         Returns a cloned network object.
-        
+
         :return: A copy of this network.
         """
         virtual_file = tempfile.SpooledTemporaryFile()
@@ -368,9 +368,9 @@ def run(
                 unclamp = unclamps.get(l, None)
                 if unclamp is not None:
                     if unclamp.ndimension() == 1:
-                        self.layers[l].s[unclamp] = 0
+                        self.layers[l].s[:, unclamp] = 0
                     else:
-                        self.layers[l].s[unclamp[t]] = 0
+                        self.layers[l].s[:, unclamp[t]] = 0
 
                 # Inject voltage to neurons.
                 inject_v = injects_v.get(l, None)
diff --git a/bindsnet/network/nodes.py b/bindsnet/network/nodes.py
index 15d4a95e..fd240de6 100644
--- a/bindsnet/network/nodes.py
+++ b/bindsnet/network/nodes.py
@@ -508,7 +508,7 @@ def forward(self, x: torch.Tensor) -> None:
 
         # Integrate inputs.
         x.masked_fill_(self.refrac_count > 0, 0.0)
-        
+
         # Decrement refractory counters.
         self.refrac_count -= self.dt
 
diff --git a/bindsnet/pipeline/environment_pipeline.py b/bindsnet/pipeline/environment_pipeline.py
index 1d954ed4..b1f50b2b 100644
--- a/bindsnet/pipeline/environment_pipeline.py
+++ b/bindsnet/pipeline/environment_pipeline.py
@@ -99,8 +99,8 @@ def __init__(
             if isinstance(layer, AbstractInput)
         ]
 
-        self.action = torch.tensor(-1)
-        self.last_action = torch.tensor(-1)
+        self.action = torch.tensor(-1, device=self.device)
+        self.last_action = torch.tensor(-1, device=self.device)
         self.action_counter = 0
         self.random_action_after = kwargs.get("random_action_after", self.time)
 
@@ -169,7 +169,7 @@ def env_step(self) -> Tuple[torch.Tensor, float, bool, Dict]:
             self.last_action = self.action
             if torch.rand(1) < self.percent_of_random_action:
                 self.action = torch.randint(
-                    low=0, high=self.spike_record[self.output].shape[-1], size=(1,)
+                    low=0, high=self.env.action_space.n, size=(1,)
                 )[0]
             elif self.action_counter > self.random_action_after:
                 if self.last_action == 0:  # last action was start b
@@ -177,7 +177,7 @@ def env_step(self) -> Tuple[torch.Tensor, float, bool, Dict]:
                     tqdm.write(f"Fire -> too many times {self.last_action} ")
                 else:
                     self.action = torch.randint(
-                        low=2, high=self.spike_record[self.output].shape[-1], size=(1,)
+                        low=0, high=self.env.action_space.n, size=(1,)
                     )[0]
                     tqdm.write(f"too many times {self.last_action} ")
             else:
diff --git a/examples/mnist/SOM_LM-SNNs.py b/examples/mnist/SOM_LM-SNNs.py
index af4b452e..2e685dfc 100644
--- a/examples/mnist/SOM_LM-SNNs.py
+++ b/examples/mnist/SOM_LM-SNNs.py
@@ -29,6 +29,7 @@
 parser.add_argument("--n_neurons", type=int, default=100)
 parser.add_argument("--n_epochs", type=int, default=1)
 parser.add_argument("--n_test", type=int, default=10000)
+parser.add_argument("--n_train", type=int, default=60000)
 parser.add_argument("--n_workers", type=int, default=-1)
 parser.add_argument("--theta_plus", type=float, default=0.05)
 parser.add_argument("--time", type=int, default=250)
@@ -48,6 +49,7 @@
 n_neurons = args.n_neurons
 n_epochs = args.n_epochs
 n_test = args.n_test
+n_train = args.n_train
 n_workers = args.n_workers
 theta_plus = args.theta_plus
 time = args.time
@@ -66,9 +68,9 @@
     torch.cuda.manual_seed_all(seed)
 else:
     torch.manual_seed(seed)
+    device = "cpu"
     if gpu:
         gpu = False
-        device = 'cpu'
 
 torch.set_num_threads(os.cpu_count() - 1)
 print("Running on Device = ", device)
@@ -106,30 +108,34 @@
 )
 
 # Record spikes during the simulation.
-spike_record = torch.zeros(update_interval, int(time/dt), n_neurons).cpu()
+spike_record = torch.zeros((update_interval, int(time / dt), n_neurons), device=device)
 
 # Neuron assignments and spike proportions.
 n_classes = 10
-assignments = -torch.ones(n_neurons).cpu()
-proportions = torch.zeros(n_neurons, n_classes).cpu()
-rates = torch.zeros(n_neurons, n_classes).cpu()
+assignments = -torch.ones(n_neurons, device=device)
+proportions = torch.zeros((n_neurons, n_classes), device=device)
+rates = torch.zeros((n_neurons, n_classes), device=device)
 
 # Sequence of accuracy estimates.
 accuracy = {"all": [], "proportion": []}
 
 # Voltage recording for excitatory and inhibitory layers.
-som_voltage_monitor = Monitor(network.layers["Y"], ["v"], time=int(time/dt))
+som_voltage_monitor = Monitor(network.layers["Y"], ["v"], time=int(time / dt))
 network.add_monitor(som_voltage_monitor, name="som_voltage")
 
 # Set up monitors for spikes and voltages
 spikes = {}
 for layer in set(network.layers):
-    spikes[layer] = Monitor(network.layers[layer], state_vars=["s"], time=int(time/dt))
+    spikes[layer] = Monitor(
+        network.layers[layer], state_vars=["s"], time=int(time / dt)
+    )
     network.add_monitor(spikes[layer], name="%s_spikes" % layer)
 
 voltages = {}
 for layer in set(network.layers) - {"X"}:
-    voltages[layer] = Monitor(network.layers[layer], state_vars=["v"], time=int(time/dt))
+    voltages[layer] = Monitor(
+        network.layers[layer], state_vars=["v"], time=int(time / dt)
+    )
     network.add_monitor(voltages[layer], name="%s_voltages" % layer)
 
 inpt_ims, inpt_axes = None, None
@@ -166,9 +172,15 @@
         dataset, batch_size=1, shuffle=True, num_workers=n_workers, pin_memory=gpu
     )
 
-    for step, batch in enumerate(tqdm(dataloader)):
+    pbar = tqdm(total=n_train)
+    for step, batch in enumerate(dataloader):
+        if step == n_train:
+            break
+
         # Get next input sample.
-        inputs = {"X": batch["encoded_image"].view(int(time/dt), 1, 1, 28, 28).to(device)}
+        inputs = {
+            "X": batch["encoded_image"].view(int(time / dt), 1, 1, 28, 28).to(device)
+        }
 
         if step > 0:
             if step % update_inhibation_weights == 0:
@@ -181,7 +193,7 @@
 
             if step % update_interval == 0:
                 # Convert the array of labels into a tensor
-                label_tensor = torch.tensor(labels).cpu()
+                label_tensor = torch.tensor(labels, device=device)
 
                 # Get network predictions.
                 all_activity_pred = all_activity(
@@ -215,7 +227,8 @@
                     )
                 )
                 tqdm.write(
-                    "Proportion weighting accuracy: %.2f (last), %.2f (average), %.2f (best)\n"
+                    "Proportion weighting accuracy: %.2f (last), %.2f (average), %.2f"
+                    " (best)\n"
                     % (
                         accuracy["proportion"][-1],
                         np.mean(accuracy["proportion"]),
@@ -247,10 +260,12 @@
             if temp_spikes.sum().sum() < 2:
                 inputs["X"] *= (
                     poisson(
-                        datum=factor * batch["image"].clamp(min=0), dt=dt, time=int(time/dt)
+                        datum=factor * batch["image"].clamp(min=0),
+                        dt=dt,
+                        time=int(time / dt),
                     )
                     .to(device)
-                    .view(int(time/dt), 1, 1, 28, 28)
+                    .view(int(time / dt), 1, 1, 28, 28)
                 )
                 factor *= factor
             else:
@@ -292,6 +307,8 @@
             plt.pause(1e-8)
 
         network.reset_state_variables()  # Reset state variables.
+        pbar.set_description_str("Train progress: ")
+        pbar.update()
 
 print("Progress: %d / %d (%.4f seconds)" % (epoch + 1, n_epochs, t() - start))
 print("Training complete.\n")
@@ -313,16 +330,19 @@
 accuracy = {"all": 0, "proportion": 0}
 
 # Record spikes during the simulation.
-spike_record = torch.zeros(1, int(time/dt), n_neurons)
+spike_record = torch.zeros(1, int(time / dt), n_neurons)
 
 # Train the network.
 print("\nBegin testing\n")
 network.train(mode=False)
 start = t()
 
-for step, batch in enumerate(tqdm(test_dataset)):
+pbar = tqdm(total=n_test)
+for step, batch in enumerate(test_dataset):
+    if step > n_test:
+        break
     # Get next input sample.
-    inputs = {"X": batch["encoded_image"].view(int(time/dt), 1, 1, 28, 28)}
+    inputs = {"X": batch["encoded_image"].view(int(time / dt), 1, 1, 28, 28)}
     if gpu:
         inputs = {k: v.cuda() for k, v in inputs.items()}
 
@@ -333,7 +353,7 @@
     spike_record[0] = spikes["Y"].get("s").squeeze()
 
     # Convert the array of labels into a tensor
-    label_tensor = torch.tensor(batch["label"])
+    label_tensor = torch.tensor(batch["label"], device=device)
 
     # Get network predictions.
     all_activity_pred = all_activity(
@@ -348,13 +368,18 @@
 
     # Compute network accuracy according to available classification strategies.
     accuracy["all"] += float(torch.sum(label_tensor.long() == all_activity_pred).item())
-    accuracy["proportion"] += float(torch.sum(label_tensor.long() == proportion_pred).item())
+    accuracy["proportion"] += float(
+        torch.sum(label_tensor.long() == proportion_pred).item()
+    )
 
     network.reset_state_variables()  # Reset state variables.
+    pbar.set_description_str("Test progress: ")
+    pbar.update()
+
 
-print("\nAll activity accuracy: %.2f" % (accuracy["all"] / test_dataset.test_labels.shape[0]))
-print("Proportion weighting accuracy: %.2f \n" % ( accuracy["proportion"] / test_dataset.test_labels.shape[0]))
+print("\nAll activity accuracy: %.2f" % (accuracy["all"] / n_test))
+print("Proportion weighting accuracy: %.2f \n" % (accuracy["proportion"] / n_test))
 
 
 print("Progress: %d / %d (%.4f seconds)" % (epoch + 1, n_epochs, t() - start))
-print("Testing complete.\n")
\ No newline at end of file
+print("Testing complete.\n")
diff --git a/examples/mnist/batch_eth_mnist.py b/examples/mnist/batch_eth_mnist.py
index a77c0cba..9f28eb9b 100644
--- a/examples/mnist/batch_eth_mnist.py
+++ b/examples/mnist/batch_eth_mnist.py
@@ -31,6 +31,7 @@
 parser.add_argument("--batch_size", type=int, default=32)
 parser.add_argument("--n_epochs", type=int, default=1)
 parser.add_argument("--n_test", type=int, default=10000)
+parser.add_argument("--n_train", type=int, default=60000)
 parser.add_argument("--n_workers", type=int, default=-1)
 parser.add_argument("--update_steps", type=int, default=256)
 parser.add_argument("--exc", type=float, default=22.5)
@@ -53,6 +54,7 @@
 batch_size = args.batch_size
 n_epochs = args.n_epochs
 n_test = args.n_test
+n_train = args.n_train
 n_workers = args.n_workers
 update_steps = args.update_steps
 exc = args.exc
@@ -68,14 +70,20 @@
 
 update_interval = update_steps * batch_size
 
+device = "cpu"
 # Sets up Gpu use
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 if gpu and torch.cuda.is_available():
     torch.cuda.manual_seed_all(seed)
 else:
     torch.manual_seed(seed)
+    device = "cpu"
     if gpu:
         gpu = False
 
+torch.set_num_threads(os.cpu_count() - 1)
+print("Running on Device = ", device)
+
 # Determines number of workers to use
 if n_workers == -1:
     n_workers = gpu * 4 * torch.cuda.device_count()
@@ -113,28 +121,32 @@
 
 # Neuron assignments and spike proportions.
 n_classes = 10
-assignments = -torch.ones(n_neurons)
-proportions = torch.zeros(n_neurons, n_classes)
-rates = torch.zeros(n_neurons, n_classes)
+assignments = -torch.ones(n_neurons, device=device)
+proportions = torch.zeros((n_neurons, n_classes), device=device)
+rates = torch.zeros((n_neurons, n_classes), device=device)
 
 # Sequence of accuracy estimates.
 accuracy = {"all": [], "proportion": []}
 
 # Voltage recording for excitatory and inhibitory layers.
-exc_voltage_monitor = Monitor(network.layers["Ae"], ["v"], time=int(time/dt))
-inh_voltage_monitor = Monitor(network.layers["Ai"], ["v"], time=int(time/dt))
+exc_voltage_monitor = Monitor(network.layers["Ae"], ["v"], time=int(time / dt))
+inh_voltage_monitor = Monitor(network.layers["Ai"], ["v"], time=int(time / dt))
 network.add_monitor(exc_voltage_monitor, name="exc_voltage")
 network.add_monitor(inh_voltage_monitor, name="inh_voltage")
 
 # Set up monitors for spikes and voltages
 spikes = {}
 for layer in set(network.layers):
-    spikes[layer] = Monitor(network.layers[layer], state_vars=["s"], time=int(time/dt))
+    spikes[layer] = Monitor(
+        network.layers[layer], state_vars=["s"], time=int(time / dt)
+    )
     network.add_monitor(spikes[layer], name="%s_spikes" % layer)
 
 voltages = {}
 for layer in set(network.layers) - {"X"}:
-    voltages[layer] = Monitor(network.layers[layer], state_vars=["v"], time=int(time/dt))
+    voltages[layer] = Monitor(
+        network.layers[layer], state_vars=["v"], time=int(time / dt)
+    )
     network.add_monitor(voltages[layer], name="%s_voltages" % layer)
 
 inpt_ims, inpt_axes = None, None
@@ -144,7 +156,7 @@
 perf_ax = None
 voltage_axes, voltage_ims = None, None
 
-spike_record = torch.zeros(update_interval, int(time/dt), n_neurons)
+spike_record = torch.zeros((update_interval, int(time / dt), n_neurons), device=device)
 
 # Train the network.
 print("\nBegin training.\n")
@@ -167,6 +179,8 @@
     )
 
     for step, batch in enumerate(tqdm(train_dataloader)):
+        if step > n_train:
+            break
         # Get next input sample.
         inputs = {"X": batch["encoded_image"]}
         if gpu:
@@ -174,7 +188,7 @@
 
         if step % update_steps == 0 and step > 0:
             # Convert the array of labels into a tensor
-            label_tensor = torch.tensor(labels)
+            label_tensor = torch.tensor(labels, device=device)
 
             # Get network predictions.
             all_activity_pred = all_activity(
@@ -208,7 +222,8 @@
                 )
             )
             print(
-                "Proportion weighting accuracy: %.2f (last), %.2f (average), %.2f (best)\n"
+                "Proportion weighting accuracy: %.2f (last), %.2f (average), %.2f"
+                " (best)\n"
                 % (
                     accuracy["proportion"][-1],
                     np.mean(accuracy["proportion"]),
@@ -263,9 +278,7 @@
             spike_ims, spike_axes = plot_spikes(spikes_, ims=spike_ims, axes=spike_axes)
             weights_im = plot_weights(square_weights, im=weights_im)
             assigns_im = plot_assignments(square_assignments, im=assigns_im)
-            perf_ax = plot_performance(
-                accuracy, ax=perf_ax
-            )
+            perf_ax = plot_performance(accuracy, ax=perf_ax)
             voltage_ims, voltage_axes = plot_voltages(
                 voltages, ims=voltage_ims, axes=voltage_axes, plot_type="line"
             )
@@ -280,8 +293,6 @@
 print("Progress: %d / %d (%.4f seconds)" % (epoch + 1, n_epochs, t() - start))
 print("Training complete.\n")
 
-
-
 # Load MNIST data.
 test_dataset = MNIST(
     PoissonEncoder(time=time, dt=dt),
@@ -311,7 +322,10 @@
 network.train(mode=False)
 start = t()
 
-for step, batch in enumerate(tqdm(test_dataloader)):
+pbar = tqdm(total=n_test)
+for step, batch in enumerate(test_dataset):
+    if step > n_test:
+        break
     # Get next input sample.
     inputs = {"X": batch["encoded_image"]}
     if gpu:
@@ -324,7 +338,7 @@
     spike_record = spikes["Ae"].get("s").permute((1, 0, 2))
 
     # Convert the array of labels into a tensor
-    label_tensor = torch.tensor(batch["label"])
+    label_tensor = torch.tensor(batch["label"], device=device)
 
     # Get network predictions.
     all_activity_pred = all_activity(
@@ -339,13 +353,16 @@
 
     # Compute network accuracy according to available classification strategies.
     accuracy["all"] += float(torch.sum(label_tensor.long() == all_activity_pred).item())
-    accuracy["proportion"] += float(torch.sum(label_tensor.long() == proportion_pred).item())
+    accuracy["proportion"] += float(
+        torch.sum(label_tensor.long() == proportion_pred).item()
+    )
 
     network.reset_state_variables()  # Reset state variables.
+    pbar.set_description_str("Test progress: ")
+    pbar.update()
 
-print("\nAll activity accuracy: %.2f" % (accuracy["all"] / test_dataset.test_labels.shape[0]))
-print("Proportion weighting accuracy: %.2f \n" % ( accuracy["proportion"] / test_dataset.test_labels.shape[0]))
-
+print("\nAll activity accuracy: %.2f" % (accuracy["all"] / n_test))
+print("Proportion weighting accuracy: %.2f \n" % (accuracy["proportion"] / n_test))
 
 print("Progress: %d / %d (%.4f seconds)" % (epoch + 1, n_epochs, t() - start))
-print("Testing complete.\n")
\ No newline at end of file
+print("Testing complete.\n")
diff --git a/examples/mnist/conv_mnist.py b/examples/mnist/conv_mnist.py
index 53ceb4f3..0a29f9bd 100644
--- a/examples/mnist/conv_mnist.py
+++ b/examples/mnist/conv_mnist.py
@@ -1,3 +1,4 @@
+import os
 import torch
 import argparse
 import matplotlib.pyplot as plt
@@ -26,6 +27,7 @@
 parser.add_argument("--seed", type=int, default=0)
 parser.add_argument("--n_epochs", type=int, default=1)
 parser.add_argument("--n_test", type=int, default=10000)
+parser.add_argument("--n_train", type=int, default=60000)
 parser.add_argument("--kernel_size", type=int, default=16)
 parser.add_argument("--stride", type=int, default=4)
 parser.add_argument("--n_filters", type=int, default=25)
@@ -46,6 +48,7 @@
 seed = args.seed
 n_epochs = args.n_epochs
 n_test = args.n_test
+n_train = args.n_train
 kernel_size = args.kernel_size
 stride = args.stride
 n_filters = args.n_filters
@@ -59,10 +62,17 @@
 plot = args.plot
 gpu = args.gpu
 
-if gpu:
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+if gpu and torch.cuda.is_available():
     torch.cuda.manual_seed_all(seed)
 else:
     torch.manual_seed(seed)
+    device = "cpu"
+    if gpu:
+        gpu = False
+
+torch.set_num_threads(os.cpu_count() - 1)
+print("Running on Device = ", device)
 
 if not train:
     update_interval = n_test
@@ -159,7 +169,8 @@
 
     for step, batch in enumerate(tqdm(train_dataloader)):
         # Get next input sample.
-
+        if step > n_train:
+            break
         inputs = {"X": batch["encoded_image"].view(time, 1, 1, 28, 28)}
         if gpu:
             inputs = {k: v.cuda() for k, v in inputs.items()}
diff --git a/examples/mnist/eth_mnist.py b/examples/mnist/eth_mnist.py
index 61720775..11771eb4 100644
--- a/examples/mnist/eth_mnist.py
+++ b/examples/mnist/eth_mnist.py
@@ -30,6 +30,7 @@
 parser.add_argument("--n_neurons", type=int, default=100)
 parser.add_argument("--n_epochs", type=int, default=1)
 parser.add_argument("--n_test", type=int, default=10000)
+parser.add_argument("--n_train", type=int, default=60000)
 parser.add_argument("--n_workers", type=int, default=-1)
 parser.add_argument("--exc", type=float, default=22.5)
 parser.add_argument("--inh", type=float, default=120)
@@ -51,6 +52,7 @@
 n_neurons = args.n_neurons
 n_epochs = args.n_epochs
 n_test = args.n_test
+n_train = args.n_train
 n_workers = args.n_workers
 exc = args.exc
 inh = args.inh
@@ -65,13 +67,18 @@
 gpu = args.gpu
 
 # Sets up Gpu use
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 if gpu and torch.cuda.is_available():
     torch.cuda.manual_seed_all(seed)
 else:
     torch.manual_seed(seed)
+    device = "cpu"
     if gpu:
         gpu = False
 
+torch.set_num_threads(os.cpu_count() - 1)
+print("Running on Device = ", device)
+
 # Determines number of workers to use
 if n_workers == -1:
     n_workers = gpu * 4 * torch.cuda.device_count()
@@ -111,32 +118,36 @@
 )
 
 # Record spikes during the simulation.
-spike_record = torch.zeros(update_interval, int(time/dt), n_neurons)
+spike_record = torch.zeros((update_interval, int(time / dt), n_neurons), device=device)
 
 # Neuron assignments and spike proportions.
 n_classes = 10
-assignments = -torch.ones(n_neurons)
-proportions = torch.zeros(n_neurons, n_classes)
-rates = torch.zeros(n_neurons, n_classes)
+assignments = -torch.ones(n_neurons, device=device)
+proportions = torch.zeros((n_neurons, n_classes), device=device)
+rates = torch.zeros((n_neurons, n_classes), device=device)
 
 # Sequence of accuracy estimates.
 accuracy = {"all": [], "proportion": []}
 
 # Voltage recording for excitatory and inhibitory layers.
-exc_voltage_monitor = Monitor(network.layers["Ae"], ["v"], time=int(time/dt))
-inh_voltage_monitor = Monitor(network.layers["Ai"], ["v"], time=int(time/dt))
+exc_voltage_monitor = Monitor(network.layers["Ae"], ["v"], time=int(time / dt))
+inh_voltage_monitor = Monitor(network.layers["Ai"], ["v"], time=int(time / dt))
 network.add_monitor(exc_voltage_monitor, name="exc_voltage")
 network.add_monitor(inh_voltage_monitor, name="inh_voltage")
 
 # Set up monitors for spikes and voltages
 spikes = {}
 for layer in set(network.layers):
-    spikes[layer] = Monitor(network.layers[layer], state_vars=["s"], time=int(time/dt))
+    spikes[layer] = Monitor(
+        network.layers[layer], state_vars=["s"], time=int(time / dt)
+    )
     network.add_monitor(spikes[layer], name="%s_spikes" % layer)
 
 voltages = {}
 for layer in set(network.layers) - {"X"}:
-    voltages[layer] = Monitor(network.layers[layer], state_vars=["v"], time=int(time/dt))
+    voltages[layer] = Monitor(
+        network.layers[layer], state_vars=["v"], time=int(time / dt)
+    )
     network.add_monitor(voltages[layer], name="%s_voltages" % layer)
 
 inpt_ims, inpt_axes = None, None
@@ -162,14 +173,16 @@
     )
 
     for step, batch in enumerate(tqdm(dataloader)):
+        if step > n_train:
+            break
         # Get next input sample.
-        inputs = {"X": batch["encoded_image"].view(int(time/dt), 1, 1, 28, 28)}
+        inputs = {"X": batch["encoded_image"].view(int(time / dt), 1, 1, 28, 28)}
         if gpu:
             inputs = {k: v.cuda() for k, v in inputs.items()}
 
         if step % update_interval == 0 and step > 0:
             # Convert the array of labels into a tensor
-            label_tensor = torch.tensor(labels)
+            label_tensor = torch.tensor(labels, device=device)
 
             # Get network predictions.
             all_activity_pred = all_activity(
@@ -203,7 +216,8 @@
                 )
             )
             print(
-                "Proportion weighting accuracy: %.2f (last), %.2f (average), %.2f (best)\n"
+                "Proportion weighting accuracy: %.2f (last), %.2f (average), %.2f"
+                " (best)\n"
                 % (
                     accuracy["proportion"][-1],
                     np.mean(accuracy["proportion"]),
@@ -279,16 +293,19 @@
 accuracy = {"all": 0, "proportion": 0}
 
 # Record spikes during the simulation.
-spike_record = torch.zeros(1, int(time/dt), n_neurons)
+spike_record = torch.zeros((1, int(time / dt), n_neurons), device=device)
 
 # Train the network.
 print("\nBegin testing\n")
 network.train(mode=False)
 start = t()
 
-for step, batch in enumerate(tqdm(test_dataset)):
+pbar = tqdm(total=n_test)
+for step, batch in enumerate(test_dataset):
+    if step > n_test:
+        break
     # Get next input sample.
-    inputs = {"X": batch["encoded_image"].view(int(time/dt), 1, 1, 28, 28)}
+    inputs = {"X": batch["encoded_image"].view(int(time / dt), 1, 1, 28, 28)}
     if gpu:
         inputs = {k: v.cuda() for k, v in inputs.items()}
 
@@ -299,7 +316,7 @@
     spike_record[0] = spikes["Ae"].get("s").squeeze()
 
     # Convert the array of labels into a tensor
-    label_tensor = torch.tensor(batch["label"])
+    label_tensor = torch.tensor(batch["label"], device=device)
 
     # Get network predictions.
     all_activity_pred = all_activity(
@@ -314,13 +331,17 @@
 
     # Compute network accuracy according to available classification strategies.
     accuracy["all"] += float(torch.sum(label_tensor.long() == all_activity_pred).item())
-    accuracy["proportion"] += float(torch.sum(label_tensor.long() == proportion_pred).item())
+    accuracy["proportion"] += float(
+        torch.sum(label_tensor.long() == proportion_pred).item()
+    )
 
     network.reset_state_variables()  # Reset state variables.
+    pbar.set_description_str("Test progress: ")
+    pbar.update()
 
-print("\nAll activity accuracy: %.2f" % (accuracy["all"] / test_dataset.test_labels.shape[0]))
-print("Proportion weighting accuracy: %.2f \n" % ( accuracy["proportion"] / test_dataset.test_labels.shape[0]))
+print("\nAll activity accuracy: %.2f" % (accuracy["all"] / n_test))
+print("Proportion weighting accuracy: %.2f \n" % (accuracy["proportion"] / n_test))
 
 
 print("Progress: %d / %d (%.4f seconds)" % (epoch + 1, n_epochs, t() - start))
-print("Testing complete.\n")
\ No newline at end of file
+print("Testing complete.\n")
diff --git a/examples/mnist/reservoir.py b/examples/mnist/reservoir.py
index a70017d1..2e2340fb 100644
--- a/examples/mnist/reservoir.py
+++ b/examples/mnist/reservoir.py
@@ -62,16 +62,16 @@
 torch.manual_seed(seed)
 
 # Sets up Gpu use
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 if gpu and torch.cuda.is_available():
-    device = "cuda"
-    torch.cuda.set_device(device)
-    # torch.set_default_tensor_type('torch.cuda.FloatTensor')
+    torch.cuda.manual_seed_all(seed)
 else:
     torch.manual_seed(seed)
     device = "cpu"
     if gpu:
         gpu = False
-
+torch.set_num_threads(os.cpu_count() - 1)
+print("Running on Device = ", device)
 
 network = Network(dt=dt)
 inpt = Input(784, shape=(1, 28, 28))
diff --git a/examples/mnist/supervised_mnist.py b/examples/mnist/supervised_mnist.py
index deac066d..3722da12 100644
--- a/examples/mnist/supervised_mnist.py
+++ b/examples/mnist/supervised_mnist.py
@@ -65,12 +65,17 @@
 device_id = args.device_id
 
 # Sets up Gpu use
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 if gpu and torch.cuda.is_available():
-    # torch.set_default_tensor_type("torch.cuda.FloatTensor")
-    torch.cuda.set_device(device_id)
     torch.cuda.manual_seed_all(seed)
 else:
     torch.manual_seed(seed)
+    device = "cpu"
+    if gpu:
+        gpu = False
+
+torch.set_num_threads(os.cpu_count() - 1)
+print("Running on Device = ", device)
 
 if not train:
     update_interval = n_test
@@ -118,18 +123,18 @@
 dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=True)
 
 # Record spikes during the simulation.
-spike_record = torch.zeros(update_interval, time, n_neurons)
+spike_record = torch.zeros(update_interval, time, n_neurons, device=device)
 
 # Neuron assignments and spike proportions.
-assignments = -torch.ones_like(torch.Tensor(n_neurons))
-proportions = torch.zeros_like(torch.Tensor(n_neurons, n_classes))
-rates = torch.zeros_like(torch.Tensor(n_neurons, n_classes))
+assignments = -torch.ones_like(torch.Tensor(n_neurons), device=device)
+proportions = torch.zeros_like(torch.Tensor(n_neurons, n_classes), device=device)
+rates = torch.zeros_like(torch.Tensor(n_neurons, n_classes), device=device)
 
 # Sequence of accuracy estimates.
 accuracy = {"all": [], "proportion": []}
 
 # Labels to determine neuron assignments and spike proportions and estimate accuracy
-labels = torch.empty(update_interval)
+labels = torch.empty(update_interval, device=device)
 
 spikes = {}
 for layer in set(network.layers):
@@ -186,7 +191,9 @@
         )
 
         # Assign labels to excitatory layer neurons.
-        assignments, proportions, rates = assign_labels(spike_record, labels, n_classes, rates)
+        assignments, proportions, rates = assign_labels(
+            spike_record, labels, n_classes, rates
+        )
 
     # Add the current label to the list of labels for this update_interval
     labels[i % update_interval] = label[0]
@@ -259,7 +266,7 @@
 accuracy = {"all": 0, "proportion": 0}
 
 # Record spikes during the simulation.
-spike_record = torch.zeros(1, int(time/dt), n_neurons)
+spike_record = torch.zeros(1, int(time / dt), n_neurons)
 
 # Train the network.
 print("\nBegin testing\n")
@@ -270,7 +277,7 @@
     if step > n_test:
         break
     # Get next input sample.
-    inputs = {"X": batch["encoded_image"].view(int(time/dt), 1, 1, 28, 28)}
+    inputs = {"X": batch["encoded_image"].view(int(time / dt), 1, 1, 28, 28)}
     if gpu:
         inputs = {k: v.cuda() for k, v in inputs.items()}
 
@@ -281,7 +288,7 @@
     spike_record[0] = spikes["Ae"].get("s").squeeze()
 
     # Convert the array of labels into a tensor
-    label_tensor = torch.tensor(batch["label"])
+    label_tensor = torch.tensor(batch["label"], device=device)
 
     # Get network predictions.
     all_activity_pred = all_activity(
@@ -296,15 +303,19 @@
 
     # Compute network accuracy according to available classification strategies.
     accuracy["all"] += float(torch.sum(label_tensor.long() == all_activity_pred).item())
-    accuracy["proportion"] += float(torch.sum(label_tensor.long() == proportion_pred).item())
+    accuracy["proportion"] += float(
+        torch.sum(label_tensor.long() == proportion_pred).item()
+    )
 
     network.reset_state_variables()  # Reset state variables.
 
-    pbar.set_description_str(f"Accuracy: {(max(accuracy['all'] ,accuracy['proportion'] ) / (step+1)):.3}")
+    pbar.set_description_str(
+        f"Accuracy: {(max(accuracy['all'] ,accuracy['proportion'] ) / (step+1)):.3}"
+    )
     pbar.update()
 
-print("\nAll activity accuracy: %.2f" % (accuracy["all"] / test_dataset.test_labels.shape[0]))
-print("Proportion weighting accuracy: %.2f \n" % (accuracy["proportion"] / test_dataset.test_labels.shape[0]))
+print("\nAll activity accuracy: %.2f" % (accuracy["all"] / n_test))
+print("Proportion weighting accuracy: %.2f \n" % (accuracy["proportion"] / n_test))
 
 
 print("Testing complete.\n")