Bug/Feature with opt O1 and ops with double

Hi, I have the following problem with opt level O1 with error message

/opt/conda/lib/python3.7/site-packages/apex/amp/wrap.py in wrapper(*args, **kwargs)
     60             raise NotImplementedError('Do not know how to handle ' +
     61                                       'these types to promote: {}'
---> 62                                       .format(types))
     63     return wrapper
     64 

NotImplementedError: Do not know how to handle these types to promote: {'DoubleTensor', 'FloatTensor'}

but the same code works with opt level O2.

How to reproduce:

import torch
from apex import amp
from torchvision.models import resnet18
model = resnet18(pretrained=False).to("cuda")
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
model, optimizer = amp.initialize(model, optimizer, opt_level="O1")
cm = torch.ones(3, 3, device='cuda').double()
cm.diag() / (cm.sum(dim=1) + cm.sum(dim=0) - cm.diag() + 1e-15)

outputs the error

Selected optimization level O1:  Insert automatic casts around Pytorch functions and Tensor methods.

Defaults for this optimization level are:
enabled                : True
opt_level              : O1
cast_model_type        : None
patch_torch_functions  : True
keep_batchnorm_fp32    : None
master_weights         : None
loss_scale             : dynamic
Processing user overrides (additional kwargs that are not None)...
After processing overrides, optimization options are:
enabled                : True
opt_level              : O1
cast_model_type        : None
patch_torch_functions  : True
keep_batchnorm_fp32    : None
master_weights         : None
loss_scale             : dynamic
---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
<ipython-input-1-851b776c17c9> in <module>
     11 cm = torch.ones(3, 3, device='cuda').double()
     12 
---> 13 cm.diag() / (cm.sum(dim=1) + cm.sum(dim=0) - cm.diag() + 1e-15)

/opt/conda/lib/python3.7/site-packages/apex/amp/wrap.py in wrapper(*args, **kwargs)
     60             raise NotImplementedError('Do not know how to handle ' +
     61                                       'these types to promote: {}'
---> 62                                       .format(types))
     63     return wrapper
     64 

NotImplementedError: Do not know how to handle these types to promote: {'DoubleTensor', 'FloatTensor'}

But the same with opt_level : O2 works

import torch
from apex import amp
from torchvision.models import resnet18
model = resnet18(pretrained=False).to("cuda")
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
model, optimizer = amp.initialize(model, optimizer, opt_level="O2")
cm = torch.ones(3, 3, device='cuda').double()
cm.diag() / (cm.sum(dim=1) + cm.sum(dim=0) - cm.diag() + 1e-15)

outputs

Selected optimization level O2:  FP16 training with FP32 batchnorm and FP32 master weights.

Defaults for this optimization level are:
enabled                : True
opt_level              : O2
cast_model_type        : torch.float16
patch_torch_functions  : False
keep_batchnorm_fp32    : True
master_weights         : True
loss_scale             : dynamic
Processing user overrides (additional kwargs that are not None)...
After processing overrides, optimization options are:
enabled                : True
opt_level              : O2
cast_model_type        : torch.float16
patch_torch_functions  : False
keep_batchnorm_fp32    : True
master_weights         : True
loss_scale             : dynamic
tensor([0.2000, 0.2000, 0.2000], device='cuda:0', dtype=torch.float64)

Any hints on how to deal with this ?
Thank you

EDIT: I solved my problem by passing to CPU

Start at apex/amp/wrap.py, specifically the wrapper line shown in the traceback, and rerun the supplied O1 reproduction. Compare its type handling with O2; completion should be a confirmed fix or an explicit documented limitation for this case.