From 5fa2a6a723de3762f1656e0536bda6795398d6e8 Mon Sep 17 00:00:00 2001 From: William Falcon Date: Tue, 14 May 2019 06:37:56 -0400 Subject: [PATCH] tng and val steps now have batch nbs --- pytorch_lightning/root_module/root_module.py | 4 ++-- setup.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/pytorch_lightning/root_module/root_module.py b/pytorch_lightning/root_module/root_module.py index f56612b9..467aeead 100644 --- a/pytorch_lightning/root_module/root_module.py +++ b/pytorch_lightning/root_module/root_module.py @@ -51,7 +51,7 @@ class RootModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks): """ raise NotImplementedError - def validation_step(self, data_batch): + def validation_step(self, data_batch, batch_nb): """ return whatever outputs will need to be aggregated in validation_end :param data_batch: @@ -67,7 +67,7 @@ class RootModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks): """ raise NotImplementedError - def training_step(self, data_batch): + def training_step(self, data_batch, batch_nb): """ return loss, dict with metrics for tqdm :param data_batch: diff --git a/setup.py b/setup.py index 94b71c20..09f5774d 100755 --- a/setup.py +++ b/setup.py @@ -7,7 +7,7 @@ from setuptools import setup, find_packages # http://blog.ionelmc.ro/2014/05/25/python-packaging/ setup( name="pytorch-lightning", - version='0.1.dev1822', + version='0.1.dev1823', description="The Keras for ML researchers using PyTorch", author="William Falcon", author_email="waf2107@columbia.edu",