Thank you @ptrblck, I thought about using forward_hook. Indeed starting from ‘O1’ we see torch.cuda.HalfTensor, but…
Here are the memory results for O0, O1 and O2 for RTX 2070 SUPER. How can we explain higher memory usage by O1? (same one can observe on google colab)
O0
check if max_memory strarts from zero:  0.0
memory after model and data are loaded:  98.30224609375
Selected optimization level O0:  Pure FP32 training.

Defaults for this optimization level are:
enabled                : True
opt_level              : O0
cast_model_type        : torch.float32
patch_torch_functions  : False
keep_batchnorm_fp32    : None
master_weights         : False
loss_scale             : 1.0
Processing user overrides (additional kwargs that are not None)...
After processing overrides, optimization options are:
enabled                : True
opt_level              : O0
cast_model_type        : torch.float32
patch_torch_functions  : False
keep_batchnorm_fp32    : None
master_weights         : False
loss_scale             : 1.0
memory after forward:  181.4619140625

O1
check if max_memory strarts from zero:  0.0
memory after model and data are loaded:  98.30224609375
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
memory after forward:  192.01171875

O2
check if max_memory strarts from zero:  0.0
memory after model and data are loaded:  98.30224609375
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
memory after forward:  98.892578125

Modified code snippet:
import torch
import torch.nn as nn

import torchvision.models as models

from apex import amp

torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

opt_level = 'O1'

print(opt_level)

seed = 123
torch.manual_seed(seed)

torch.cuda.reset_max_memory_allocated()

print('check if max_memory strarts from zero: ', torch.cuda.max_memory_allocated()/1024**2)

model = models.resnet50()
model.cuda()

#model.avgpool.register_forward_hook(lambda m, x, y: print(x[0].type(), y[0].type()))
x = torch.randn(1, 3, 224, 224).cuda()

torch.cuda.synchronize()
print('memory after model and data are loaded: ', torch.cuda.max_memory_allocated()/1024**2)

model = amp.initialize(model, opt_level=opt_level)
output = model(x)

torch.cuda.synchronize()
print('memory after forward: ', torch.cuda.max_memory_allocated()/1024**2)