How can I add input caption, especially for image?

[mrzResearchArena]

I need to add an image caption. Could you please assist me to add it?

[pranavkantgaur]

@mrzResearchArena The issue with using \caption and also \name command is that they are also defined in this package, see here:
https://github.com/HarisIqbal88/PlotNeuralNet/blob/master/layers/Ball.sty#L28
https://github.com/HarisIqbal88/PlotNeuralNet/blob/master/layers/Ball.sty#L29

---

I resolved the issue by renaming these commands.

[mrzResearchArena]

Thanks for your response. I tried to add a caption in the image input. However, I could not do it yet.
In addition, it will be helpful for me if you assist more about it.

https://pastebin.com/NExDJAyY

[Sirius291]

Same problem here.
My workaround is this: I generate the tikzpicture depicting my architecture separately, as a single file.
I compile with pdflatex and treat the resulting pdf as an image in my actual latex document. When you use the package "graphicx", you can simply add pdfs as an image via the "\includegraphics" command and put a figure environment around it as usual.