mirror of
https://github.com/jantic/DeOldify.git
synced 2026-10-31 03:14:33 +08:00
Gareth Davidson
dff82b6d2e
All the files marked as executable were bugging me, so I ran this: find . -executable -type f -print0 | xargs -0 grep -L '#!' | xargs chmod -x
13 lines
454 B
Python
13 lines
454 B
Python
" Miscellaneous callbacks "
|
|
|
|
from fastai.callback import Callback
|
|
|
|
class StopAfterNBatches(Callback):
|
|
"Stop training after n batches of the first epoch."
|
|
def __init__(self, n_batches:int=2):
|
|
self.stop,self.n_batches = False,n_batches-1 # iteration starts from 0
|
|
|
|
def on_batch_end(self, iteration, **kwargs):
|
|
if iteration == self.n_batches:
|
|
return {'stop_epoch': True, 'stop_training': True, 'skip_validate': True}
|