Plotly를 이용한 multiple images 출력

import plotly.express as px
from plotly.subplots import make_subplots
import plotly.graph_objects as go

# Create a 2x2 subplot layout
fig = make_subplots(rows=2, cols=2)

# Add each image as a heatmap
fig.add_trace(go.Heatmap(z=label.squeeze().detach().to('cpu').numpy(),
                         colorscale='gray', showscale=False), row=1, col=1)
fig.add_trace(go.Heatmap(z=prediction.squeeze().detach().to('cpu').numpy(),
                         colorscale='gray', showscale=False), row=1, col=2)
fig.add_trace(go.Heatmap(z=error_mag.squeeze().detach().to('cpu').numpy(),
                         colorscale='gray', showscale=False), row=2, col=1)
fig.add_trace(go.Heatmap(z=mask.squeeze().detach().to('cpu').numpy(),
                         colorscale='gray', showscale=False), row=2, col=2)

fig.update_layout(height=600, width=600, title_text="4 Images in One Figure")
fig.show()