RuntimeError: expected mat1 and mat2 to have the same dtype, but got: float != c10::BFloat16

Describe the bug
I am using deepspeed with PyTorch Lightning and during training, it always fails on the nn.TransformerLayerEncoder. In my deepspeed config, I set bf16 to true, I tried autocast, and manually converting the params to bfloat16 , but unfortunately it always fails with matmul dtype mismatch. Are transformer layers compatible with deepspeed? Seems like deepspeed is having trouble converting the weights to bfloat16. They still appear to be float32

Deepspeed Config:

zero_allow_untested_optimizer: true
bf16:
  enabled: true
fp16:
  enabled: false
activation_checkpointing:
  partition_activations: true
  cpu_checkpointing: true
  contiguous_memory_optimization: true
  synchronize_checkpoint_boundary: true
  profile: false
zero_optimization:
  stage: 2
  offload_optimizer:
    device: cpu
    pin_memory: true
  offload_param:
    device: cpu
    pin_memory: true
  overlap_comm: true
  contiguous_gradients: true
  reduce_bucket_size: "auto"
  reduce_scatter: true
  allgather_partitions: true
  allgather_bucket_size: "auto"
  reduce_bucket_size: "auto"
gradient_clipping: "auto"

transformer_encoder = nn.TransformerEncoder(
                nn.TransformerEncoderLayer(
                    d_model=input_dim,
                    nhead=num_heads,
                    dim_feedforward=mlp_dim,
                    dropout=dropout,
                    activation=F.silu,
                    batch_first=True,
                    # norm_first=True,
                ),
                num_layers=num_transformers,
            )