Getting error when running python code on tensorflow-metal on Macbook pro M1

I have successfully installed tensorflow-metal on my Macbook pro M1 in a separate environment, but when I run the following sample code:

import tensorflow as tf
from tensorflow.keras import datasets, layers, models

(train_images, train_labels), (test_images, test_labels) = datasets.mnist.load_data()
train_images, test_images = train_images / 255.0, test_images / 255.0

model = models.Sequential([
    layers.Flatten(input_shape=(28, 28)),
    layers.Dense(128, activation='relu'),
    layers.Dense(64, activation='relu'),
    layers.Dense(10)
])
model.compile(
    optimizer='adam',
    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=['accuracy']
)
history = model.fit(
    train_images, 
    train_labels, 
    epochs=10, 
    validation_data=(test_images, test_labels)
)

I got these errors:

2021-10-10 09:26:08.398847: I
tensorflow/compiler/mlir/mlir_graph_optimization_pass.cc:185] None of
the MLIR Optimization Passes are enabled (registered 2) 2021-10-10
09:26:08.399120: W
tensorflow/core/platform/profile_utils/cpu_utils.cc:128] Failed to get
CPU frequency: 0 Hz Epoch 1/10 2021-10-10 09:26:08.554082: I
tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.cc:112]
Plugin optimizer for device_type GPU is enabled. 2021-10-10
09:26:08.581 python[65936:4662219] -[MPSGraph
adamUpdateWithLearningRateTensor:beta1Tensor:beta2Tensor:epsilonTensor:beta1PowerTensor:beta2PowerTensor:valuesTensor:momentumTensor:velocityTensor:maximumVelocityTensor:gradientTensor:name:]:
unrecognized selector sent to instance 0x16a877460 2021-10-10
09:26:08.594 python[65936:4662219] *** Terminating app due to uncaught
exception ‘NSInvalidArgumentException’, reason: ‘-[MPSGraph
adamUpdateWithLearningRateTensor:beta1Tensor:beta2Tensor:epsilonTensor:beta1PowerTensor:beta2PowerTensor:valuesTensor:momentumTensor:velocityTensor:maximumVelocityTensor:gradientTensor:name:]:
unrecognized selector sent to instance 0x16a877460’
*** First throw call stack: ( 0 CoreFoundation 0x00000001929a5320 __exceptionPreprocess + 240 1 libobjc.A.dylib 0x00000001926d3c04 objc_exception_throw + 60 2 CoreFoundation 0x0000000192a34020
-[NSObject(NSObject) __retain_OA] + 0 3 CoreFoundation 0x0000000192907174 forwarding + 1444 4 CoreFoundation
0x0000000192906b20 _CF_forwarding_prep_0 + 96 5
libmetal_plugin.dylib 0x000000012e81a290
_ZN12metal_plugin14MPSApplyAdamOpIfEC2EPNS_20OpKernelConstructionE + 656 6 libmetal_plugin.dylib 0x000000012e819ebc
_ZN12metal_pluginL14CreateOpKernelINS_14MPSApplyAdamOpIfEEEEPvP23TF_OpKernelConstruction

• 52 7 libtensorflow_framework.2.dylib 0x00000001192545d4 _ZN10tensorflow12_GLOBAL__N_120KernelBuilderFactory6CreateEPNS_20OpKernelConstructionE
• 88 8 libtensorflow_framework.2.dylib 0x00000001192d6158 _ZN10tensorflow14CreateOpKernelENS_10DeviceTypeEPNS_10DeviceBaseEPNS_9AllocatorEPNS_22FunctionLibraryRuntimeEPNS_11ResourceMgrERKNSt3__110shared_ptrIKNS_14NodePropertiesEEEiPPNS_8OpKernelE
• 784 9 libtensorflow_framework.2.dylib 0x00000001194b12b8 _ZN10tensorflow21CreateNonCachedKernelEPNS_6DeviceEPNS_22FunctionLibraryRuntimeERKNSt3__110shared_ptrIKNS_14NodePropertiesEEEiPPNS_8OpKernelE
• 272 10 libtensorflow_framework.2.dylib 0x000000011945bc20 _ZN10tensorflow26FunctionLibraryRuntimeImpl12CreateKernelERKNSt3__110shared_ptrIKNS_14NodePropertiesEEEPNS_22FunctionLibraryRuntimeEPPNS_8OpKernelE
• 600 11 libtensorflow_framework.2.dylib 0x00000001194c6430 _ZN10tensorflow22ImmutableExecutorState10InitializeERKNS_5GraphE + 1192 12 libtensorflow_framework.2.dylib 0x00000001194b1064
  _ZN10tensorflow16NewLocalExecutorERKNS_19LocalExecutorParamsERKNS_5GraphEPPNS_8ExecutorE
• 304 13 libtensorflow_framework.2.dylib 0x00000001194bee6c _ZN10tensorflow12_GLOBAL__N_124DefaultExecutorRegistrar7Factory11NewExecutorERKNS_19LocalExecutorParamsERKNS_5GraphEPNSt3__110unique_ptrINS_8ExecutorENS9_14default_deleteISB_EEEE + 48 14 libtensorflow_framework.2.dylib 0x00000001194bf7e8 _ZN10tensorflow11NewExecutorERKNSt3__112basic_stringIcNS0_11char_traitsIcEENS0_9allocatorIcEEEERKNS_19LocalExecutorParamsERKNS_5GraphEPNS0_10unique_ptrINS_8ExecutorENS0_14default_deleteISG_EEEE
• 92 15 libtensorflow_framework.2.dylib 0x000000011945e278 _ZN10tensorflow26FunctionLibraryRuntimeImpl10CreateItemEPPNS0_4ItemE + 2676 16 libtensorflow_framework.2.dylib 0x000000011945f06c
  _ZN10tensorflow26FunctionLibraryRuntimeImpl3RunERKNS_22FunctionLibraryRuntime7OptionsEyN4absl12lts_202103244SpanIKNS_6TensorEEEPNSt3__16vectorIS8_NSB_9allocatorIS8_EEEENSB_8functionIFvRKNS_6StatusEEEE
• 676 17 libtensorflow_framework.2.dylib 0x000000011946d0c0 _ZNK10tensorflow29ProcessFunctionLibraryRuntime14RunMultiDeviceERKNS_22FunctionLibraryRuntime7OptionsEyPNSt3__16vectorIN4absl12lts_202103247variantIJNS_6TensorENS_11TensorShapeEEEENS5_9allocatorISC_EEEEPNS6_INS5_10unique_ptrINS0_11CleanUpItemENS5_14default_deleteISI_EEEENSD_ISL_EEEENS5_8functionIFvRKNS_6StatusEEEENSP_IFSQ_RKNS0_21ComponentFunctionDataEPNS0_12InternalArgsEEEE
• 2640 18 libtensorflow_framework.2.dylib 0x0000000119470098 _ZNK10tensorflow29ProcessFunctionLibraryRuntime3RunERKNS_22FunctionLibraryRuntime7OptionsEyN4absl12lts_202103244SpanIKNS_6TensorEEEPNSt3__16vectorIS8_NSB_9allocatorIS8_EEEENSB_8functionIFvRKNS_6StatusEEEE
• 2012 19 libtensorflow_framework.2.dylib 0x0000000119470868 _ZNK10tensorflow29ProcessFunctionLibraryRuntime7RunSyncERKNS_22FunctionLibraryRuntime7OptionsEyN4absl12lts_202103244SpanIKNS_6TensorEEEPNSt3__16vectorIS8_NSB_9allocatorIS8_EEEE + 160 20 _pywrap_tensorflow_internal.so 0x000000014c175538 _ZN10tensorflow19KernelAndDeviceFunc3RunEPNS_19ScopedStepContainerERKNS_15EagerKernelArgsEPNSt3__16vectorIN4absl12lts_202103247variantIJNS_6TensorENS_11TensorShapeEEEENS6_9allocatorISD_EEEEPNS_19CancellationManagerERKNS9_8optionalINS_25EagerRemoteFunctionParamsEEERKNSK_INS_17ManagedStackTraceEEE
• 516 21 _pywrap_tensorflow_internal.so 0x000000014c13fd44 _ZN10tensorflow18EagerKernelExecuteEPNS_12EagerContextERKN4absl12lts_2021032413InlinedVectorIPNS_12TensorHandleELm4ENSt3__19allocatorIS6_EEEERKNS3_8optionalINS_25EagerRemoteFunctionParamsEEERKNS7_10unique_ptrINS_15KernelAndDeviceENS_4core15RefCountDeleterEEEPNS_14GraphCollectorEPNS_19CancellationManagerENS3_4SpanIS6_EERKNSD_INS_17ManagedStackTraceEEE + 372 22 _pywrap_tensorflow_internal.so 0x000000014c1463a8 _ZN10tensorflow11ExecuteNode3RunEv + 396 23 _pywrap_tensorflow_internal.so 0x000000014c481748 _ZN10tensorflow13EagerExecutor11SyncExecuteEPNS_9EagerNodeE + 172 24 _pywrap_tensorflow_internal.so 0x000000014c13f880 _ZN10tensorflow12_GLOBAL__N_117EagerLocalExecuteEPNS_14EagerOperationEPPNS_12TensorHandleEPi
• 1976 25 _pywrap_tensorflow_internal.so 0x000000014c13da28 _ZN10tensorflow12EagerExecuteEPNS_14EagerOperationEPPNS_12TensorHandleEPi
• 296 26 _pywrap_tensorflow_internal.so 0x000000014bda2b88 _ZN10tensorflow14EagerOperation7ExecuteEN4absl12lts_202103244SpanIPNS_20AbstractTensorHandleEEEPi
• 192 27 _pywrap_tensorflow_internal.so 0x000000014c17b910 _ZN10tensorflow21CustomDeviceOpHandler7ExecuteEPNS_27ImmediateExecutionOperationEPPNS_30ImmediateExecutionTensorHandleEPi
• 468 28 _pywrap_tensorflow_internal.so 0x00000001489c7f1c TFE_Execute + 80 29 _pywrap_tensorflow_internal.so
  0x0000000148944934
  _Z24TFE_Py_ExecuteCancelableP11TFE_ContextPKcS2_PN4absl12lts_2021032413InlinedVectorIP16TFE_TensorHandleLm4ENSt3__19allocatorIS7_EEEEP7_objectP23TFE_CancellationManagerPNS5_IS7_Lm2ESA_EEP9TF_Status
• 616 30 _pywrap_tfe.so 0x000000011a6ea4c0 ZN10tensorflow32TFE_Py_ExecuteCancelable_wrapperERKN8pybind116handleEPKcS5_S3_S3_PNS_19CancellationManagerES3 + 160 31 _pywrap_tfe.so 0x000000011a71b220 _ZZN8pybind1112cpp_function10initializeIZL25pybind11_init__pywrap_tfeRNS_7module_EE4$44NS_6objectEJRKNS_6handleEPKcSA_S8_S8_S8_EJNS_4nameENS_5scopeENS_7siblingEEEEvOT_PFT0_DpT1_EDpRKT2_ENUlRNS_6detail13function_callEE_8__invokeESR
• 184 32 _pywrap_tfe.so 0x000000011a6fd16c ZN8pybind1112cpp_function10dispatcherEP7_objectS2_S2 + 3216 33 python 0x00000001043e58cc
  cfunction_call_varargs + 292 34 python
  0x00000001043e4f50 _PyObject_MakeTpCall + 640 35 python
  0x00000001044fa3e0 call_function + 680 36 python
  0x00000001044f6ebc _PyEval_EvalFrameDefault + 29472 37 python
  0x00000001044ef584 _PyEval_EvalCodeWithName + 3340 38 python
  0x00000001043e5f10 _PyFunction_Vectorcall + 236 39 python
  0x00000001044fa348 call_function + 528 40 python
  0x00000001044f6fac _PyEval_EvalFrameDefault + 29712 41 python
  0x00000001044ef584 _PyEval_EvalCodeWithName + 3340 42 python
  0x00000001043e5f10 _PyFunction_Vectorcall + 236 43 python
  0x00000001043e991c method_vectorcall + 156 44 python
  0x00000001044fa348 call_function + 528 45 python
  0x00000001044f6fac _PyEval_EvalFrameDefault + 29712 46 python
  0x00000001044ef584 _PyEval_EvalCodeWithName + 3340 47 python
  0x00000001043e5f10 _PyFunction_Vectorcall + 236 48 python
  0x00000001043e991c method_vectorcall + 156 49 python
  0x00000001044fa348 call_function + 528 50 python
  0x00000001044f6fac _PyEval_EvalFrameDefault + 29712 51 python
  0x00000001044ef584 _PyEval_EvalCodeWithName + 3340 52 python
  0x00000001043e5f10 _PyFunction_Vectorcall + 236 53 python
  0x00000001043e4b84 _PyObject_FastCallDict + 132 54 python
  0x00000001043e6cc4 _PyObject_Call_Prepend + 156 55 python
  0x000000010445a774 slot_tp_call + 296 56 python
  0x00000001043e56e8 PyObject_Call + 312 57 python
  0x00000001044f714c _PyEval_EvalFrameDefault + 30128 58 python
  0x00000001044ef584 _PyEval_EvalCodeWithName + 3340 59 python
  0x00000001043e5f10 _PyFunction_Vectorcall + 236 60 python
  0x00000001043e9a88 method_vectorcall + 520 61 python
  0x00000001043e54c0 PyVectorcall_Call + 120 62 python
  0x00000001044f714c _PyEval_EvalFrameDefault + 30128 63 python
  0x00000001044ef584 _PyEval_EvalCodeWithName + 3340 64 python
  0x00000001043e5f10 _PyFunction_Vectorcall + 236 65 python
  0x00000001043e4c1c _PyObject_FastCallDict + 284 66 python
  0x00000001043e6cc4 _PyObject_Call_Prepend + 156 67 python
  0x000000010445a774 slot_tp_call + 296 68 python
  0x00000001043e4f50 _PyObject_MakeTpCall + 640 69 python
  0x00000001044fa3e0 call_function + 680 70 python
  0x00000001044f6ebc _PyEval_EvalFrameDefault + 29472 71 python
  0x00000001044ef584 _PyEval_EvalCodeWithName + 3340 72 python
  0x00000001043e5f10 _PyFunction_Vectorcall + 236 73 python
  0x00000001043e991c method_vectorcall + 156 74 python
  0x00000001044fa348 call_function + 528 75 python
  0x00000001044f6fac _PyEval_EvalFrameDefault + 29712 76 python
  0x00000001044ef584 _PyEval_EvalCodeWithName + 3340 77 python
  0x000000010454c0b0 PyRun_InteractiveOneObjectEx + 1196 78 python
  0x00000001045488d0 PyRun_InteractiveLoopFlags + 540 79 python
  0x000000010454863c PyRun_AnyFileExFlags + 72 80 python
  0x00000001045710fc Py_RunMain + 3568 81 python
  0x00000001045720a8 pymain_main + 1244 82 python
  0x00000001043bc1e0 main + 56 83 libdyld.dylib
  0x0000000192848f54 start + 4 ) libc++abi.dylib: terminating with
  uncaught exception of type NSException zsh: abort python

I cannot fix this error, please help me. Thank you so much.