title: Python numpy simple way to filter label zero from train dataset
body: <p>I have a train dataset trainX and trainy.
  TrainX shape is (n,128,3) trainy shape is (n,1).</p>
  <p>I would like to remove all samples where the trainy label is zero. I have the following code which is half working, but some of the dimensions of the numpy arrays are lost.</p>
  <pre><code># load the dataset, returns train and test X and y elements
  def load_dataset_train_only(prefix='trainingData128/'):
      # load all train
      trainX, trainy = load_dataset_group('train', prefix)            
      print(trainX.shape, trainy.shape)
      # zero-offset class values
      trainy = trainy
  
      #remove samples where stroke rate = 0
      train_filter = np.where(trainy != 0)
      trainX, trainy = trainX[train_filter], trainy[train_filter]
      print('Size after filtering off zero labels',trainX.shape, trainy.shape)
              
      # one hot encode y
      trainy = to_categorical(trainy)
      print(trainX.shape, trainy.shape)
      return trainX, trainy
  </code></pre>
  <p>The shape before filtering is;
  (366511, 128, 3) (366511, 1)</p>
  <p>And the shape after filtering is;
  Size after filtering off zero labels (280905, 3) (280905,)</p>
  <p>The 280905 value is correct, so it is filtering the zero labels, but how do I modify the code so I don't loose the dimensions?</p>
  <p>Thanks</p>