@pytest.mark.push
    def test_mnist_training():
        torch.manual_seed(0)
    
        # Model and data type.
        # For bfloat16, the following line should be added to the test_forge_vs_torch function:
        # In file forge/forge/op/eval/forge/eltwise_unary.py:418 should be replaced with: threshold_tensor = ac.tensor(torch.zeros(shape, dtype=torch.bfloat16) + threshold)
        # That sets relu threshold to bfloat16 tensor.
        # And in file forge/forge/compile.py::compile_main forced bfloat 16 should be added compiler_cfg.default_df_override = DataFormat.Float16_b
        dtype = torch.float32
    
        # Set training hyperparameters
        num_epochs = 3
        batch_size = 2048
        learning_rate = 0.001
    
        # Load dataset
        test_loader, train_loader = load_dataset(batch_size, dtype=dtype)
    
        # Define model and instruct it to compile and run on TT device
        framework_model = MNISTLinear(
            bias=False, dtype=dtype
        )  # bias=False because batch_size=1 with bias=True is not supported
    
        # Create a torch loss and leave on CPU
        loss_fn = torch.nn.CrossEntropyLoss()
    
        # Define optimizer and instruct it to compile and run on TT device
        framework_optimizer = torch.optim.SGD(framework_model.parameters(), lr=learning_rate)
        tt_model = forge.compile(
            framework_model,
            sample_inputs=[torch.rand(batch_size, 784, dtype=dtype)],
            optimizer=framework_optimizer,
            training=True,
        )
    
        logger.info("Starting training loop... (logger will be disabled)")
        logger.disable("")
        for epoch_idx in range(num_epochs):
    
            total_loss = 0
            for batch_idx, (data, target) in enumerate(train_loader):
                # Reset gradients (every batch)
                framework_optimizer.zero_grad()
    
                # Create target tensor and leave on CPU
                target = nn.functional.one_hot(target, num_classes=10).to(dtype)
    
                # Forward pass (prediction) on device
>               pred = tt_model(data)[0]

forge/test/mlir/mnist/training/test_training.py:68: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = <forge.compiled_graph_state.CompiledModel object at 0x7fd7abb72a70>
inputs = (tensor([[-0.42421, -0.42421, -0.42421,  ..., -0.42421, -0.42421, -0.42421],
        [-0.42421, -0.42421, -0.42421,  ....21,  ..., -0.42421, -0.42421, -0.42421],
        [-0.42421, -0.42421, -0.42421,  ..., -0.42421, -0.42421, -0.42421]]),)
inputs_and_parameters = [tensor([[-0.42421, -0.42421, -0.42421,  ..., -0.42421, -0.42421, -0.42421],
        [-0.42421, -0.42421, -0.42421,  ....  0.02915,  0.01170],
        [ 0.00670,  0.03483, -0.01020,  ..., -0.03968, -0.00866, -0.00669]], requires_grad=True)]
name = 'linear_relu_stack.4.weight'
param = Parameter containing:
tensor([[ 0.02984, -0.01652,  0.04047,  ..., -0.03471, -0.02169, -0.00426],
        [ 0.02000, -...,  0.02915,  0.01170],
        [ 0.00670,  0.03483, -0.01020,  ..., -0.03968, -0.00866, -0.00669]], requires_grad=True)

    def __call__(self, *inputs: AnyTensor) -> List[torch.Tensor]:
        """
        Run inference on the compiled model.
    
        Parameters
        ----------
        inputs: [Tensor, ...]
            Input tensors
    
        Returns
        -------
        List[Tensor]
            Output tensors
        """
        self.inputs = [*inputs]
        inputs_and_parameters = [
            *inputs,
            *self.fwd_compiled_graph_state.get_ordered_constant_tensors(),
            *self.fwd_compiled_graph_state.get_ordered_parameter_tensors(),
        ]
    
        if any([not isinstance(t, torch.Tensor) for t in inputs_and_parameters]):
            logger.info("Converting inputs and parameters to PyTorch tensors...")
            inputs_and_parameters = to_pt_tensors(inputs_and_parameters)
    
        if self.training() and isinstance(self.framework_module, PyTorchModule):
            for name, param in self.framework_module.module.named_parameters():
                if param.requires_grad:
                    our_tensor = self.fwd_compiled_graph_state.get_parameter_tensor(name)
    
                    # NOTE: for parameters that require gradients, we want to share the same tensor with the PyTorch module.
                    # This is because we want to be able to optimize the parameters both on the device (through our runtime)
                    # and via the torch optimizers. So this ensures that whichever side updates the parameter value, the other side can see the change.
                    #
                    # This could change in the future, but for now ensure that our premise is correct.
                    assert param is our_tensor
    
        logger.info(
            f"Running model {self.framework_module.get_name()} {self.fwd_compiled_graph_state.graph.get_name()} on device..."
        )
>       all_outputs = run_binary(self.compiled_binary, int(ProgramId.FORWARD), inputs_and_parameters)
E       RuntimeError: TT_THROW @ /__w/tt-forge-fe/tt-forge-fe/third_party/tt-mlir/third_party/tt-metal/src/tt-metal/tt_metal/impl/program/program.cpp:819: tt::exception
E       info:
E       Statically allocated circular buffers on core range [(x=0,y=0) - (x=3,y=0)] grow to 1671968 B which is beyond max L1 size of 1499136 B
E       backtrace:
E        --- /__w/tt-forge-fe/tt-forge-fe/third_party/tt-mlir/third_party/tt-metal/src/tt-metal-build/lib/libtt_metal.so(+0x14d99b) [0x7fd808d0699b]
E        --- tt::tt_metal::detail::Program_::validate_circular_buffer_region(tt::tt_metal::v0::Device const*)
E        --- tt::tt_metal::v0::CommandQueue::run_command_impl(tt::tt_metal::CommandInterface const&)
E        --- tt::tt_metal::v0::EnqueueProgram(tt::tt_metal::v0::CommandQueue&, tt::tt_metal::v0::Program&, bool)
E        --- void ttnn::device_operation::detail::launch_on_worker_thread<tt::tt_metal::operation::OldInfraDeviceOperation<std::vector<tt::tt_metal::Tensor, std::allocator<tt::tt_metal::Tensor> > >, unsigned char, long, tt::tt_metal::operation::DeviceOperation<std::vector<tt::tt_metal::Tensor, std::allocator<tt::tt_metal::Tensor> > >, tt::tt_metal::operation::OldInfraDeviceOperation<std::vector<tt::tt_metal::Tensor, std::allocator<tt::tt_metal::Tensor> > >::tensor_args_t, std::vector<tt::tt_metal::Tensor, std::allocator<tt::tt_metal::Tensor> >, tt::tt_metal::v0::Device*>(unsigned char, long, tt::tt_metal::operation::DeviceOperation<std::vector<tt::tt_metal::Tensor, std::allocator<tt::tt_metal::Tensor> > > const&, tt::tt_metal::operation::OldInfraDeviceOperation<std::vector<tt::tt_metal::Tensor, std::allocator<tt::tt_metal::Tensor> > >::tensor_args_t const&, std::vector<tt::tt_metal::Tensor, std::allocator<tt::tt_metal::Tensor> >&, tt::tt_metal::v0::Device*&)
E        --- tt::tt_metal::operation::OldInfraDeviceOperation<std::vector<tt::tt_metal::Tensor, std::allocator<tt::tt_metal::Tensor> > >::tensor_return_value_t ttnn::device_operation::detail::launch_on_single_device<tt::tt_metal::operation::OldInfraDeviceOperation<std::vector<tt::tt_metal::Tensor, std::allocator<tt::tt_metal::Tensor> > > >(unsigned char, tt::tt_metal::operation::OldInfraDeviceOperation<std::vector<tt::tt_metal::Tensor, std::allocator<tt::tt_metal::Tensor> > >::operation_attributes_t const&, tt::tt_metal::operation::OldInfraDeviceOperation<std::vector<tt::tt_metal::Tensor, std::allocator<tt::tt_metal::Tensor> > >::tensor_args_t const&)
E        --- /__w/tt-forge-fe/tt-forge-fe/third_party/tt-mlir/third_party/tt-metal/src/tt-metal-build/lib/_ttnn.so(+0x1e15769) [0x7fd7f77c7769]
E        --- tt::tt_metal::operation::OldInfraDeviceOperation<std::vector<tt::tt_metal::Tensor, std::allocator<tt::tt_metal::Tensor> > >::tensor_return_value_t ttnn::device_operation::detail::invoke<tt::tt_metal::operation::OldInfraDeviceOperation<std::vector<tt::tt_metal::Tensor, std::allocator<tt::tt_metal::Tensor> > > >(unsigned char, tt::tt_metal::operation::OldInfraDeviceOperation<std::vector<tt::tt_metal::Tensor, std::allocator<tt::tt_metal::Tensor> > >::operation_attributes_t const&, tt::tt_metal::operation::OldInfraDeviceOperation<std::vector<tt::tt_metal::Tensor, std::allocator<tt::tt_metal::Tensor> > >::tensor_args_t const&)
E        --- /__w/tt-forge-fe/tt-forge-fe/third_party/tt-mlir/third_party/tt-metal/src/tt-metal-build/lib/_ttnn.so(+0x1e1208f) [0x7fd7f77c408f]
E        --- std::vector<tt::tt_metal::Tensor, std::allocator<tt::tt_metal::Tensor> > tt::tt_metal::operation::run<std::vector<tt::tt_metal::Tensor, std::allocator<tt::tt_metal::Tensor> > >(tt::tt_metal::operation::DeviceOperation<std::vector<tt::tt_metal::Tensor, std::allocator<tt::tt_metal::Tensor> > >&&, std::vector<tt::tt_metal::Tensor, std::allocator<tt::tt_metal::Tensor> > const&, std::vector<std::optional<tt::tt_metal::Tensor const>, std::allocator<std::optional<tt::tt_metal::Tensor const> > > const&, std::vector<std::optional<tt::tt_metal::Tensor>, std::allocator<std::optional<tt::tt_metal::Tensor> > > const&, unsigned char)
E        --- /__w/tt-forge-fe/tt-forge-fe/third_party/tt-mlir/third_party/tt-metal/src/tt-metal-build/lib/_ttnn.so(+0x1429f03) [0x7fd7f6ddbf03]
E        --- /__w/tt-forge-fe/tt-forge-fe/third_party/tt-mlir/third_party/tt-metal/src/tt-metal-build/lib/_ttnn.so(+0x142a9c5) [0x7fd7f6ddc9c5]
E        --- ttnn::operations::matmul::matmul(tt::tt_metal::Tensor const&, tt::tt_metal::Tensor const&, std::optional<tt::tt_metal::Tensor const> const&, ttnn::operations::matmul::Matmul const&, unsigned char)
E        --- ttnn::operations::matmul::bound_matmul(tt::tt_metal::Tensor const&, tt::tt_metal::Tensor const&, std::optional<tt::tt_metal::Tensor const> const&, ttnn::operations::matmul::Matmul const&, unsigned char const&)
E        --- /__w/tt-forge-fe/tt-forge-fe/third_party/tt-mlir/third_party/tt-metal/src/tt-metal-build/lib/_ttnn.so(_ZN4ttnn10operations6matmul15MatmulOperation6invokeERKN2tt8tt_metal6TensorES7_bbRKSt8optionalIKNS4_12MemoryConfigEES8_IKNS4_8DataTypeEERKS8_IKSt7variantIJNS1_28MatmulMultiCoreProgramConfigENS1_45MatmulMultiCoreNonOptimizedReuseProgramConfigENS1_33MatmulMultiCoreReuseProgramConfigENS1_42MatmulMultiCoreReuseMultiCastProgramConfigENS1_44MatmulMultiCoreReuseMultiCast1DProgramConfigENS1_53MatmulMultiCoreReuseMultiCastDRAMShardedProgramConfigEEEERKS8_IKNSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEEEES8_IKSH_IJNS_28GrayskullComputeKernelConfigENS_27WormholeComputeKernelConfigEEEES8_IKNS_5types8CoreGridEERKS8_IKNS4_4TileEE+0x2c5) [0x7fd7f6d732c5]
E        --- /__w/tt-forge-fe/tt-forge-fe/third_party/tt-mlir/build/lib/SharedLib/libTTMLIR.so(+0x7784db) [0x7fd809e704db]
E        --- /__w/tt-forge-fe/tt-forge-fe/third_party/tt-mlir/build/lib/SharedLib/libTTMLIR.so(+0x7777ad) [0x7fd809e6f7ad]
E        --- tt::runtime::ttnn::operations::matmul::run(tt::target::ttnn::MatmulOp const*, tt::runtime::ttnn::ProgramContext&)
E        --- tt::runtime::ttnn::runProgram(tt::tt_metal::distributed::MeshDevice&, tt::runtime::Binary, unsigned int, std::vector<tt::tt_metal::Tensor*, std::allocator<tt::tt_metal::Tensor*> > const&)
E        --- tt::runtime::ttnn::submit(tt::runtime::Device, tt::runtime::Binary, unsigned int, std::vector<tt::runtime::Tensor, std::allocator<tt::runtime::Tensor> > const&)
E        --- tt::runtime::submit(tt::runtime::Device, tt::runtime::Binary, unsigned int, std::vector<tt::runtime::Tensor, std::allocator<tt::runtime::Tensor> > const&)
E        --- tt::run_binary(tt::runtime::Binary&, int, std::vector<at::Tensor, std::allocator<at::Tensor> > const&)
E        --- /__w/tt-forge-fe/tt-forge-fe/forge/forge/_C.so(+0x310600) [0x7fd80a3d5600]
E        --- /__w/tt-forge-fe/tt-forge-fe/forge/forge/_C.so(+0x31054e) [0x7fd80a3d554e]
E        --- /__w/tt-forge-fe/tt-forge-fe/forge/forge/_C.so(+0xdf8f5) [0x7fd80a1a48f5]
E        --- /opt/ttforge-toolchain/venv/bin/python(+0x18b282) [0x559b0b1c1282]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyObject_MakeTpCall+0x25b) [0x559b0b1b7b4b]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyEval_EvalFrameDefault+0x5c7d) [0x559b0b1b0f4d]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyObject_FastCallDictTstate+0xc4) [0x559b0b1b6d14]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyObject_Call_Prepend+0x5c) [0x559b0b1cbd4c]
E        --- /opt/ttforge-toolchain/venv/bin/python(+0x29e054) [0x559b0b2d4054]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyObject_MakeTpCall+0x25b) [0x559b0b1b7b4b]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyEval_EvalFrameDefault+0x5c7d) [0x559b0b1b0f4d]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyFunction_Vectorcall+0x7c) [0x559b0b1c1aec]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyEval_EvalFrameDefault+0x2c89) [0x559b0b1adf59]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyFunction_Vectorcall+0x7c) [0x559b0b1c1aec]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyEval_EvalFrameDefault+0x2c89) [0x559b0b1adf59]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyFunction_Vectorcall+0x7c) [0x559b0b1c1aec]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyEval_EvalFrameDefault+0x1a22) [0x559b0b1accf2]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyFunction_Vectorcall+0x7c) [0x559b0b1c1aec]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyEval_EvalFrameDefault+0x58aa) [0x559b0b1b0b7a]
E        --- /opt/ttforge-toolchain/venv/bin/python(+0x198be1) [0x559b0b1cebe1]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyEval_EvalFrameDefault+0x58aa) [0x559b0b1b0b7a]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyFunction_Vectorcall+0x7c) [0x559b0b1c1aec]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyObject_FastCallDictTstate+0x16d) [0x559b0b1b6dbd]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyObject_Call_Prepend+0x5c) [0x559b0b1cbd4c]
E        --- /opt/ttforge-toolchain/venv/bin/python(+0x29e054) [0x559b0b2d4054]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyObject_MakeTpCall+0x25b) [0x559b0b1b7b4b]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyEval_EvalFrameDefault+0x671a) [0x559b0b1b19ea]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyFunction_Vectorcall+0x7c) [0x559b0b1c1aec]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyEval_EvalFrameDefault+0x818) [0x559b0b1abae8]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyFunction_Vectorcall+0x7c) [0x559b0b1c1aec]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyEval_EvalFrameDefault+0x2c89) [0x559b0b1adf59]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyFunction_Vectorcall+0x7c) [0x559b0b1c1aec]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyEval_EvalFrameDefault+0x1a22) [0x559b0b1accf2]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyFunction_Vectorcall+0x7c) [0x559b0b1c1aec]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyEval_EvalFrameDefault+0x58aa) [0x559b0b1b0b7a]
E        --- /opt/ttforge-toolchain/venv/bin/python(+0x198be1) [0x559b0b1cebe1]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyEval_EvalFrameDefault+0x58aa) [0x559b0b1b0b7a]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyFunction_Vectorcall+0x7c) [0x559b0b1c1aec]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyObject_FastCallDictTstate+0x16d) [0x559b0b1b6dbd]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyObject_Call_Prepend+0x5c) [0x559b0b1cbd4c]
E        --- /opt/ttforge-toolchain/venv/bin/python(+0x29e054) [0x559b0b2d4054]
E        --- /opt/ttforge-toolchain/venv/bin/python(PyObject_Call+0xbb) [0x559b0b1cf81b]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyEval_EvalFrameDefault+0x2c89) [0x559b0b1adf59]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyFunction_Vectorcall+0x7c) [0x559b0b1c1aec]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyEval_EvalFrameDefault+0x6d2) [0x559b0b1ab9a2]
E        --- /opt/ttforge-toolchain/venv/bin/python(+0x198be1) [0x559b0b1cebe1]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyEval_EvalFrameDefault+0x1a22) [0x559b0b1accf2]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyFunction_Vectorcall+0x7c) [0x559b0b1c1aec]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyEval_EvalFrameDefault+0x2c89) [0x559b0b1adf59]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyFunction_Vectorcall+0x7c) [0x559b0b1c1aec]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyEval_EvalFrameDefault+0x6d2) [0x559b0b1ab9a2]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyFunction_Vectorcall+0x7c) [0x559b0b1c1aec]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyEval_EvalFrameDefault+0x1a22) [0x559b0b1accf2]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyFunction_Vectorcall+0x7c) [0x559b0b1c1aec]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyEval_EvalFrameDefault+0x2c89) [0x559b0b1adf59]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyFunction_Vectorcall+0x7c) [0x559b0b1c1aec]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyEval_EvalFrameDefault+0x1a22) [0x559b0b1accf2]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyFunction_Vectorcall+0x7c) [0x559b0b1c1aec]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyEval_EvalFrameDefault+0x58aa) [0x559b0b1b0b7a]
E        --- /opt/ttforge-toolchain/venv/bin/python(+0x198be1) [0x559b0b1cebe1]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyEval_EvalFrameDefault+0x58aa) [0x559b0b1b0b7a]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyFunction_Vectorcall+0x7c) [0x559b0b1c1aec]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyObject_FastCallDictTstate+0x16d) [0x559b0b1b6dbd]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyObject_Call_Prepend+0x5c) [0x559b0b1cbd4c]
E        --- /opt/ttforge-toolchain/venv/bin/python(+0x29e054) [0x559b0b2d4054]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyObject_MakeTpCall+0x25b) [0x559b0b1b7b4b]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyEval_EvalFrameDefault+0x671a) [0x559b0b1b19ea]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyFunction_Vectorcall+0x7c) [0x559b0b1c1aec]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyEval_EvalFrameDefault+0x2c89) [0x559b0b1adf59]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyFunction_Vectorcall+0x7c) [0x559b0b1c1aec]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyEval_EvalFrameDefault+0x1a22) [0x559b0b1accf2]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyFunction_Vectorcall+0x7c) [0x559b0b1c1aec]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyEval_EvalFrameDefault+0x58aa) [0x559b0b1b0b7a]
E        --- /opt/ttforge-toolchain/venv/bin/python(+0x198be1) [0x559b0b1cebe1]
E        --- /opt/ttforge-toolchain/venv/bin/python(_PyEval_EvalFrameDefault+0x58aa) [0x559b0b1b0b7a]

forge/forge/compiled_graph_state.py:273: RuntimeError