added option to change default tensor

This commit is contained in:
William Falcon 2019-05-13 21:22:50 -04:00
parent 38717abcd4
commit 4dbf38093a
1 changed files with 1 additions and 1 deletions

View File

@ -219,13 +219,13 @@ class Trainer(TrainerIO):
# filter out the weights that were done on gpu so we can load on good old cpus
self.optimizers = model.configure_optimizers()
pdb.set_trace()
if self.use_amp:
# An example
self.model, optimizer = amp.initialize(
self.model, self.optimizers[0], opt_level="O2",
keep_batchnorm_fp32=True, loss_scale="dynamic"
)
pdb.set_trace()
self.optimizers[0] = optimizer
model.trainer = self