Issue with multi gpu training 
I get the following error on training with multi gpus:
  File "training_on_each_split.py", line 271, in <module>
    trainer.fit(model)
  File "/home/nvarshn2/.conda/envs/pytorch_lightning/lib/python3.6/site-packages/pytorch_lightning/trainer/trainer.py", line 997, in fit
    results = self.dp_train(model)
  File "/home/nvarshn2/.conda/envs/pytorch_lightning/lib/python3.6/site-packages/pytorch_lightning/trainer/distrib_parts.py", line 270, in dp_train
    result = self.run_pretrain_routine(model)
  File "/home/nvarshn2/.conda/envs/pytorch_lightning/lib/python3.6/site-packages/pytorch_lightning/trainer/trainer.py", line 1213, in run_pretrain_routine
    self.train()
  File "/home/nvarshn2/.conda/envs/pytorch_lightning/lib/python3.6/site-packages/pytorch_lightning/trainer/training_loop.py", line 370, in train
    self.run_training_epoch()
  File "/home/nvarshn2/.conda/envs/pytorch_lightning/lib/python3.6/site-packages/pytorch_lightning/trainer/training_loop.py", line 452, in run_training_epoch
    batch_output = self.run_training_batch(batch, batch_idx)
  File "/home/nvarshn2/.conda/envs/pytorch_lightning/lib/python3.6/site-packages/pytorch_lightning/trainer/training_loop.py", line 632, in run_training_batch
    self.hiddens
  File "/home/nvarshn2/.conda/envs/pytorch_lightning/lib/python3.6/site-packages/pytorch_lightning/trainer/training_loop.py", line 783, in optimizer_closure
    training_step_output = self.process_output(training_step_output, train=True)
  File "/home/nvarshn2/.conda/envs/pytorch_lightning/lib/python3.6/site-packages/pytorch_lightning/trainer/logging.py", line 112, in process_output
    callback_metrics = self.reduce_distributed_output(callback_metrics, num_gpus)
  File "/home/nvarshn2/.conda/envs/pytorch_lightning/lib/python3.6/site-packages/pytorch_lightning/trainer/logging.py", line 193, in reduce_distributed_output
    output[k] = self.reduce_distributed_output(output[k], num_gpus)
  File "/home/nvarshn2/.conda/envs/pytorch_lightning/lib/python3.6/site-packages/pytorch_lightning/trainer/logging.py", line 201, in reduce_distributed_output
    output[k] = torch.mean(output[k])
RuntimeError: Can only calculate the mean of floating types. Got Long instead.


Code
    def loss_function(self, logits, labels):

        with torch.no_grad():
            predicted_label = logits.max(dim=1)[1]
            prediction_label_count = {}
            for label in range(parameters["num_labels"]):
                prediction_label_count[label] = predicted_label.eq(label).sum()

        loss_fn = torch.nn.CrossEntropyLoss()
        eturn loss_fn(logits, labels), prediction_label_count

    def training_step(self, batch, batch_idx):
        logits, softmax_logits = self.forward(**batch)
        loss, prediction_label_count = self.loss_function(logits, batch["labels"])
        return {"loss": loss,
                "prediction_label_count": prediction_label_count,
                "log": {"train_loss": loss}}

Need help in debugging the issue.