diff --git a/openwakeword/data.py b/openwakeword/data.py
index c91d7f3..964ce77 100755
--- a/openwakeword/data.py
+++ b/openwakeword/data.py
@@ -451,7 +451,7 @@ def mix_clips_batch(
         # Apply volume augmentation
         if volume_augmentation:
             volume_levels = np.random.uniform(0.02, 1.0, mixed_clips_batch.shape[0])
-            mixed_clips_batch = (volume_levels/mixed_clips_batch.max(axis=1)[0])[..., None]*mixed_clips_batch
+            mixed_clips_batch = (volume_levels/mixed_clips_batch.max(dim=1)[0])[..., None]*mixed_clips_batch
         else:
             # Normalize clips only if max value is outside of [-1, 1]
             abs_max, _ = torch.max(
@@ -463,7 +463,7 @@ def mix_clips_batch(
         mixed_clips_batch = (mixed_clips_batch.numpy()*32767).astype(np.int16)
 
         # Remove any clips that are silent (happens rarely when mixing/reverberating)
-        error_index = np.where(mixed_clips_batch.max(axis=1) != 0)[0]
+        error_index = torch.from_numpy(np.where(mixed_clips_batch.max(dim=1) != 0)[0])
         mixed_clips_batch = mixed_clips_batch[error_index]
         labels_batch = labels_batch[error_index]
         sequence_labels_batch = sequence_labels_batch[error_index]
@@ -686,7 +686,7 @@ def augment_clips(
 
         # Do second pass augmentations
         device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
-        augmented_batch = augment2(samples=torch.vstack(augmented_clips).unsqueeze(axis=1).to(device), sample_rate=sr).squeeze(axis=1)
+        augmented_batch = augment2(samples=torch.vstack(augmented_clips).unsqueeze(dim=1).to(device), sample_rate=sr).squeeze(axis=1)
 
         # Do reverberation
         if augmentation_probabilities["RIR"] >= np.random.random() and RIR_paths != []:
diff --git a/openwakeword/train.py b/openwakeword/train.py
index 9bb280f..b741627 100755
--- a/openwakeword/train.py
+++ b/openwakeword/train.py
@@ -410,6 +410,7 @@ class Model(nn.Module):
             if step_ndx == max_steps-1:
                 break
 
+
 # Separate function to convert onnx models to tflite format
 def convert_onnx_to_tflite(onnx_model_path, output_path):
     """Converts an ONNX version of an openwakeword model to the Tensorflow tflite format."""
@@ -697,14 +698,15 @@ if __name__ == '__main__':
             batch_size=len(X_val_fp_labels)
         )
 
-        X_val = np.vstack((
-            np.load(os.path.join(feature_save_dir, "positive_features_test.npy")),
-            np.load(os.path.join(feature_save_dir, "negative_features_test.npy"))
-        ))
-        labels = np.hstack((np.ones(X_val.shape[0]//2), np.zeros(X_val.shape[0]//2))).astype(np.float32)
+        X_val_pos = np.load(os.path.join(feature_save_dir, "positive_features_test.npy"))
+        X_val_neg = np.load(os.path.join(feature_save_dir, "negative_features_test.npy"))
+        labels = np.hstack((np.ones(X_val_pos.shape[0]), np.zeros(X_val_neg.shape[0]))).astype(np.float32)
 
         X_val = torch.utils.data.DataLoader(
-            torch.utils.data.TensorDataset(torch.from_numpy(X_val), torch.from_numpy(labels)),
+            torch.utils.data.TensorDataset(
+                torch.from_numpy(np.vstack((X_val_pos, X_val_neg))),
+                torch.from_numpy(labels)
+                ),
             batch_size=len(labels)
         )