modified for SD Layer

python/paddle/nn/layer/common.py

@@ -171,6 +171,11 @@ class Linear(Layer):
         self.name = name
 
     def forward(self, input):
+        with paddle.amp.auto_cast(custom_white_list={'elementwise_add','fused_gemm_epilogue'}, dtype='bfloat16'):
+            out = paddle.incubate.nn.functional.fused_linear(
+                x=input, weight=self.weight, bias=self.bias, name=self.name
+            )
+            return out
         out = F.linear(
             x=input, weight=self.weight, bias=self.bias, name=self.name
         )

python/paddle/nn/layer/conv.py

@@ -15,7 +15,7 @@
 # TODO: define classes of convolutional neural network
 
 import numpy as np
-
+import paddle
 from paddle import get_flags
 
 from ...device import (
@@ -704,6 +704,7 @@ class Conv2D(_ConvNd):
         )
 
     def forward(self, x):
+        with paddle.amp.auto_cast(custom_white_list={'elementwise_add'}, level='O1', dtype='bfloat16'):
             if self._padding_mode != 'zeros':
                 x = F.pad(
                     x,