I think that the signature of _board is a bit too specific. And the way it is called in the trainer class definitely is too specific: it assumes that there are always two elements in batch, which might not always be the case (think of a variational auto encoder).
I think that, at the Module level, _board should have a signature that looks more like that of forward. Although we will also give it the output of the forward pass, so maybe two arguments which are tuples:
class Module:
...
def _board(self, tb, inputs, outputs):
pass
I think that the signature of
_board
is a bit too specific. And the way it is called in the trainer class definitely is too specific: it assumes that there are always two elements inbatch
, which might not always be the case (think of a variational auto encoder).I think that, at the
Module
level,_board
should have a signature that looks more like that offorward
. Although we will also give it the output of the forward pass, so maybe two arguments which are tuples:Then in SegNet knows what to expect
And in the trainer you can have:
Finally, I think that the name can be "public" (
board
instead of_board
).What do you think? It akes the signature of
board
a bit ugly but at least it is generic.