@user_123454321
Thanks for reply. I get the following when I print model definition:
==========================================================================================
Layer (type:depth-idx)                   Output Shape              Param #
├─Sequential: 1-1                        [-1, 16, 512, 512]        –
|    └─ReflectionPad2d: 2-1              [-1, 3, 516, 516]         –
|    └─Conv2d: 2-2                       [-1, 16, 512, 512]        1,216
|    └─SELU: 2-3                         [-1, 16, 512, 512]        –
|    └─BatchNorm2d: 2-4                  [-1, 16, 512, 512]        32
├─Sequential: 1-2                        [-1, 32, 256, 256]        –
|    └─ReflectionPad2d: 2-5              [-1, 16, 516, 516]        –
|    └─Conv2d: 2-6                       [-1, 32, 256, 256]        12,832
|    └─SELU: 2-7                         [-1, 32, 256, 256]        –
|    └─BatchNorm2d: 2-8                  [-1, 32, 256, 256]        64
├─Sequential: 1-3                        [-1, 64, 128, 128]        –
|    └─ReflectionPad2d: 2-9              [-1, 32, 260, 260]        –
|    └─Conv2d: 2-10                      [-1, 64, 128, 128]        51,264
|    └─SELU: 2-11                        [-1, 64, 128, 128]        –
|    └─BatchNorm2d: 2-12                 [-1, 64, 128, 128]        128
├─Sequential: 1-4                        [-1, 128, 64, 64]         –
|    └─ReflectionPad2d: 2-13             [-1, 64, 132, 132]        –
|    └─Conv2d: 2-14                      [-1, 128, 64, 64]         204,928
|    └─SELU: 2-15                        [-1, 128, 64, 64]         –
|    └─BatchNorm2d: 2-16                 [-1, 128, 64, 64]         256
├─Sequential: 1-5                        [-1, 128, 32, 32]         –
|    └─ReflectionPad2d: 2-17             [-1, 128, 68, 68]         –
|    └─Conv2d: 2-18                      [-1, 128, 32, 32]         409,728
|    └─SELU: 2-19                        [-1, 128, 32, 32]         –
|    └─BatchNorm2d: 2-20                 [-1, 128, 32, 32]         256
├─Sequential: 1-6                        [-1, 128, 16, 16]         –
|    └─ReflectionPad2d: 2-21             [-1, 128, 36, 36]         –
|    └─Conv2d: 2-22                      [-1, 128, 16, 16]         409,728
├─Sequential: 1-7                        [-1, 128, 8, 8]           –
|    └─ReflectionPad2d: 2-23             [-1, 128, 20, 20]         –
|    └─Conv2d: 2-24                      [-1, 128, 8, 8]           409,728
├─Conv2d: 1-8                            [-1, 128, 1, 1]           1,048,704
├─Sequential: 1-9                        [-1, 128, 1, 1]           –
|    └─SELU: 2-25                        [-1, 128, 1, 1]           –
|    └─Conv2d: 2-26                      [-1, 128, 1, 1]           16,512
├─Sequential: 1-10                       [-1, 128, 32, 32]         –
|    └─Conv2d: 2-27                      [-1, 128, 32, 32]         16,512
Total params: 2,581,888
Trainable params: 2,581,888
Non-trainable params: 0
Total mult-adds (G): 3.40